Add integration test for UpdateDatasource API

Signed-off-by: Heemin Kim <heemin@amazon.com>

src/main/java/org/opensearch/geospatial/exceptions/IncompatibleDatasourceException.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.geospatial.exceptions;
+
+import java.io.IOException;
+
+import org.opensearch.OpenSearchException;
+import org.opensearch.common.io.stream.StreamInput;
+import org.opensearch.rest.RestStatus;
+
+/**
+ * IncompatibleDatasourceException corresponding to the {@link RestStatus#BAD_REQUEST} status code
+ *
+ * The exception is thrown when a user tries to update datasource with new endpoint which is not compatible
+ * with current datasource
+ */
+public class IncompatibleDatasourceException extends OpenSearchException {
+
+    public IncompatibleDatasourceException(String msg, Object... args) {
+        super(msg, args);
+    }
+
+    public IncompatibleDatasourceException(String msg, Throwable cause, Object... args) {
+        super(msg, cause, args);
+    }
+
+    public IncompatibleDatasourceException(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    public final RestStatus status() {
+        return RestStatus.BAD_REQUEST;
+    }
+}

src/main/java/org/opensearch/geospatial/ip2geo/action/RestUpdateDatasourceHandler.java

@@ -31,15 +31,15 @@ public String getName() {
 
     @Override
     protected RestChannelConsumer prepareRequest(final RestRequest request, final NodeClient client) throws IOException {
-        final PutDatasourceRequest putDatasourceRequest = new PutDatasourceRequest(request.param("name"));
+        final UpdateDatasourceRequest updateDatasourceRequest = new UpdateDatasourceRequest(request.param("name"));
         if (request.hasContentOrSourceParam()) {
             try (XContentParser parser = request.contentOrSourceParamParser()) {
-                PutDatasourceRequest.PARSER.parse(parser, putDatasourceRequest, null);
+                UpdateDatasourceRequest.PARSER.parse(parser, updateDatasourceRequest, null);
             }
         }
         return channel -> client.executeLocally(
             UpdateDatasourceAction.INSTANCE,
-            putDatasourceRequest,
+            updateDatasourceRequest,
             new RestToXContentListener<>(channel)
         );
     }

src/main/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceRequest.java

@@ -33,8 +33,8 @@
 @Log4j2
 @EqualsAndHashCode(callSuper = false)
 public class UpdateDatasourceRequest extends AcknowledgedRequest<UpdateDatasourceRequest> {
-    private static final ParseField ENDPOINT_FIELD = new ParseField("endpoint");
-    private static final ParseField UPDATE_INTERVAL_IN_DAYS_FIELD = new ParseField("update_interval_in_days");
+    public static final ParseField ENDPOINT_FIELD = new ParseField("endpoint");
+    public static final ParseField UPDATE_INTERVAL_IN_DAYS_FIELD = new ParseField("update_interval_in_days");
     private static final int MAX_DATASOURCE_NAME_BYTES = 255;
     /**
      * @param name the datasource name

src/main/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceTransportAction.java

@@ -21,6 +21,7 @@
 import org.opensearch.action.support.HandledTransportAction;
 import org.opensearch.action.support.master.AcknowledgedResponse;
 import org.opensearch.common.inject.Inject;
+import org.opensearch.geospatial.exceptions.IncompatibleDatasourceException;
 import org.opensearch.geospatial.ip2geo.common.DatasourceFacade;
 import org.opensearch.geospatial.ip2geo.common.DatasourceManifest;
 import org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService;
@@ -146,7 +147,7 @@ private void validateFieldsCompatibility(final UpdateDatasourceRequest request,
 
         List<String> fields = datasourceUpdateService.getHeaderFields(request.getEndpoint());
         if (datasource.isCompatible(fields) == false) {
-            throw new OpenSearchException(
+            throw new IncompatibleDatasourceException(
                 "new fields [{}] does not contain all old fields [{}]",
                 fields.toString(),
                 datasource.getDatabase().getFields().toString()

src/test/java/org/opensearch/geospatial/GeospatialRestTestCase.java

@@ -110,9 +110,9 @@ protected static void deletePipeline(String name) throws IOException {
         client().performRequest(request);
     }
 
-    protected Response createDatasource(final String name, Map<String, String> properties) throws IOException {
+    protected Response createDatasource(final String name, Map<String, Object> properties) throws IOException {
         XContentBuilder builder = XContentFactory.jsonBuilder().startObject();
-        for (Map.Entry<String, String> config : properties.entrySet()) {
+        for (Map.Entry<String, Object> config : properties.entrySet()) {
             builder.field(config.getKey(), config.getValue());
         }
         builder.endObject();
@@ -149,16 +149,31 @@ protected Response deleteDatasource(final String name) throws IOException {
         return client().performRequest(request);
     }
 
+    protected Response deleteDatasource(final String name, final int retry) throws Exception {
+        for (int i = 0; i < retry; i++) {
+            try {
+                Request request = new Request(DELETE, buildDatasourcePath(name));
+                return client().performRequest(request);
+            } catch (Exception e) {
+                if (i + 1 == retry) {
+                    throw e;
+                }
+                Thread.sleep(1000);
+            }
+        }
+        throw new RuntimeException("should not reach here");
+    }
+
     protected Map<String, Object> getDatasource(final String name) throws Exception {
         Request request = new Request(GET, buildDatasourcePath(name));
         Response response = client().performRequest(request);
         return createParser(XContentType.JSON.xContent(), EntityUtils.toString(response.getEntity())).map();
     }
 
-    protected Response updateDatasource(final String name, Map<String, Object> config) throws IOException {
+    protected Response updateDatasource(final String name, Map<String, Object> properties) throws IOException {
         XContentBuilder builder = XContentFactory.jsonBuilder().startObject();
-        if (config != null && !config.isEmpty()) {
-            builder.value(config);
+        for (Map.Entry<String, Object> config : properties.entrySet()) {
+            builder.field(config.getKey(), config.getValue());
         }
         builder.endObject();
 

src/test/java/org/opensearch/geospatial/exceptions/IncompatibleDatasourceExceptionTests.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.geospatial.exceptions;
+
+import org.opensearch.rest.RestStatus;
+import org.opensearch.test.OpenSearchTestCase;
+
+public class IncompatibleDatasourceExceptionTests extends OpenSearchTestCase {
+    public void testStatusCode_whenCreated_thenBadRequestStatus() {
+        IncompatibleDatasourceException exception = new IncompatibleDatasourceException(
+            "New datasource is not compatible with existing datasource"
+        );
+        assertEquals(RestStatus.BAD_REQUEST, exception.status());
+
+        IncompatibleDatasourceException exceptionWithCause = new IncompatibleDatasourceException(
+            "New datasource is not compatible with existing datasource",
+            new RuntimeException()
+        );
+        assertEquals(RestStatus.BAD_REQUEST, exceptionWithCause.status());
+    }
+}

src/test/java/org/opensearch/geospatial/exceptions/ResourceInUseExceptionTests.java

@@ -0,0 +1,16 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.geospatial.exceptions;
+
+import org.opensearch.rest.RestStatus;
+import org.opensearch.test.OpenSearchTestCase;
+
+public class ResourceInUseExceptionTests extends OpenSearchTestCase {
+    public void testStatusCode() {
+        ResourceInUseException exception = new ResourceInUseException("Resource is in use");
+        assertEquals(RestStatus.BAD_REQUEST, exception.status());
+    }
+}

src/test/java/org/opensearch/geospatial/ip2geo/Ip2GeoTestCase.java

@@ -183,7 +183,7 @@ protected Datasource randomDatasource() {
         Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS);
         Datasource datasource = new Datasource();
         datasource.setName(GeospatialTestHelper.randomLowerCaseString());
-        datasource.setSchedule(new IntervalSchedule(now, Randomness.get().nextInt(29), ChronoUnit.DAYS));
+        datasource.setSchedule(new IntervalSchedule(now, Randomness.get().nextInt(28) + 1, ChronoUnit.DAYS));
         datasource.setState(randomState());
         datasource.setIndices(Arrays.asList(GeospatialTestHelper.randomLowerCaseString(), GeospatialTestHelper.randomLowerCaseString()));
         datasource.setEndpoint(String.format(Locale.ROOT, "https://%s.com/manifest.json", GeospatialTestHelper.randomLowerCaseString()));

src/test/java/org/opensearch/geospatial/ip2geo/action/RestUpdateDatasourceHandlerTests.java

@@ -41,11 +41,34 @@ public void testPrepareRequest_whenValidInput_thenSucceed() {
         AtomicBoolean isExecuted = new AtomicBoolean(false);
 
         verifyingClient.setExecuteLocallyVerifier((actionResponse, actionRequest) -> {
-            assertTrue(actionRequest instanceof PutDatasourceRequest);
-            PutDatasourceRequest putDatasourceRequest = (PutDatasourceRequest) actionRequest;
-            assertEquals("https://test.com", putDatasourceRequest.getEndpoint());
-            assertEquals(TimeValue.timeValueDays(1), putDatasourceRequest.getUpdateInterval());
-            assertEquals(datasourceName, putDatasourceRequest.getName());
+            assertTrue(actionRequest instanceof UpdateDatasourceRequest);
+            UpdateDatasourceRequest updateDatasourceRequest = (UpdateDatasourceRequest) actionRequest;
+            assertEquals("https://test.com", updateDatasourceRequest.getEndpoint());
+            assertEquals(TimeValue.timeValueDays(1), updateDatasourceRequest.getUpdateInterval());
+            assertEquals(datasourceName, updateDatasourceRequest.getName());
+            isExecuted.set(true);
+            return null;
+        });
+
+        dispatchRequest(request);
+        assertTrue(isExecuted.get());
+    }
+
+    public void testPrepareRequest_whenNullInput_thenSucceed() {
+        String datasourceName = GeospatialTestHelper.randomLowerCaseString();
+        String content = "{}";
+        RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.PUT)
+            .withPath(String.format(Locale.ROOT, path, datasourceName))
+            .withContent(new BytesArray(content), XContentType.JSON)
+            .build();
+        AtomicBoolean isExecuted = new AtomicBoolean(false);
+
+        verifyingClient.setExecuteLocallyVerifier((actionResponse, actionRequest) -> {
+            assertTrue(actionRequest instanceof UpdateDatasourceRequest);
+            UpdateDatasourceRequest updateDatasourceRequest = (UpdateDatasourceRequest) actionRequest;
+            assertNull(updateDatasourceRequest.getEndpoint());
+            assertNull(updateDatasourceRequest.getUpdateInterval());
+            assertEquals(datasourceName, updateDatasourceRequest.getName());
             isExecuted.set(true);
             return null;
         });

src/test/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceIT.java

@@ -0,0 +1,104 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.geospatial.ip2geo.action;
+
+import java.io.IOException;
+import java.time.Duration;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+
+import lombok.SneakyThrows;
+
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.opensearch.client.ResponseException;
+import org.opensearch.geospatial.GeospatialRestTestCase;
+import org.opensearch.geospatial.GeospatialTestHelper;
+import org.opensearch.geospatial.ip2geo.Ip2GeoDataServer;
+import org.opensearch.rest.RestStatus;
+
+public class UpdateDatasourceIT extends GeospatialRestTestCase {
+    // Use this value in resource name to avoid name conflict among tests
+    private static final String PREFIX = UpdateDatasourceIT.class.getSimpleName().toLowerCase(Locale.ROOT);
+
+    @BeforeClass
+    public static void start() {
+        Ip2GeoDataServer.start();
+    }
+
+    @AfterClass
+    public static void stop() {
+        Ip2GeoDataServer.stop();
+    }
+
+    @SneakyThrows
+    public void testUpdateDatasource_whenValidInput_thenUpdated() {
+        boolean isDatasourceCreated = false;
+        String datasourceName = PREFIX + GeospatialTestHelper.randomLowerCaseString();
+        try {
+            Map<String, Object> datasourceProperties = Map.of(
+                PutDatasourceRequest.ENDPOINT_FIELD.getPreferredName(),
+                Ip2GeoDataServer.getEndpointCountry()
+            );
+
+            // Create datasource and wait for it to be available
+            createDatasource(datasourceName, datasourceProperties);
+            isDatasourceCreated = true;
+            waitForDatasourceToBeAvailable(datasourceName, Duration.ofSeconds(10));
+
+            int updateIntervalInDays = 1;
+            updateDatasourceEndpoint(datasourceName, Ip2GeoDataServer.getEndpointCity(), updateIntervalInDays);
+            List<Map<String, Object>> datasources = (List<Map<String, Object>>) getDatasource(datasourceName).get("datasources");
+
+            assertEquals(Ip2GeoDataServer.getEndpointCity(), datasources.get(0).get("endpoint"));
+            assertEquals(updateIntervalInDays, datasources.get(0).get("update_interval_in_days"));
+        } finally {
+            if (isDatasourceCreated) {
+                deleteDatasource(datasourceName, 3);
+            }
+        }
+    }
+
+    @SneakyThrows
+    public void testUpdateDatasource_whenIncompatibleFields_thenFails() {
+        boolean isDatasourceCreated = false;
+        String datasourceName = PREFIX + GeospatialTestHelper.randomLowerCaseString();
+        try {
+            Map<String, Object> datasourceProperties = Map.of(
+                PutDatasourceRequest.ENDPOINT_FIELD.getPreferredName(),
+                Ip2GeoDataServer.getEndpointCity()
+            );
+
+            // Create datasource and wait for it to be available
+            createDatasource(datasourceName, datasourceProperties);
+            isDatasourceCreated = true;
+            waitForDatasourceToBeAvailable(datasourceName, Duration.ofSeconds(10));
+
+            // Update should fail as country data does not have every fields that city data has
+            int updateIntervalInDays = 1;
+            ResponseException exception = expectThrows(
+                ResponseException.class,
+                () -> updateDatasourceEndpoint(datasourceName, Ip2GeoDataServer.getEndpointCountry(), updateIntervalInDays)
+            );
+            assertEquals(RestStatus.BAD_REQUEST.getStatus(), exception.getResponse().getStatusLine().getStatusCode());
+        } finally {
+            if (isDatasourceCreated) {
+                deleteDatasource(datasourceName, 3);
+            }
+        }
+    }
+
+    private void updateDatasourceEndpoint(final String datasourceName, final String endpoint, final int updateInterval) throws IOException {
+        Map<String, Object> properties = Map.of(
+            UpdateDatasourceRequest.ENDPOINT_FIELD.getPreferredName(),
+            endpoint,
+            UpdateDatasourceRequest.UPDATE_INTERVAL_IN_DAYS_FIELD.getPreferredName(),
+            updateInterval
+        );
+        updateDatasource(datasourceName, properties);
+    }
+}

src/test/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceTransportActionTests.java

@@ -27,6 +27,7 @@
 import org.opensearch.action.support.master.AcknowledgedResponse;
 import org.opensearch.common.Randomness;
 import org.opensearch.common.unit.TimeValue;
+import org.opensearch.geospatial.exceptions.IncompatibleDatasourceException;
 import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase;
 import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource;
 import org.opensearch.jobscheduler.spi.LockModel;
@@ -197,7 +198,7 @@ public void testDoExecute_whenIncompatibleFields_thenError() {
         // Verify
         ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(Exception.class);
         verify(listener).onFailure(exceptionCaptor.capture());
-        assertEquals(OpenSearchException.class, exceptionCaptor.getValue().getClass());
+        assertEquals(IncompatibleDatasourceException.class, exceptionCaptor.getValue().getClass());
         exceptionCaptor.getValue().getMessage().contains("does not contain");
         verify(ip2GeoLockService).releaseLock(eq(lockModel), any(ActionListener.class));
     }