Adding a cleanup method to EnterpriseGeoIpDownloaderIT

Author: masseyke

modules/ingest-geoip/src/internalClusterTest/java/org/elasticsearch/ingest/geoip/EnterpriseGeoIpDownloaderIT.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.ResourceAlreadyExistsException;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.LatchedActionListener;
 import org.elasticsearch.action.bulk.BulkItemResponse;
 import org.elasticsearch.action.bulk.BulkRequest;
 import org.elasticsearch.action.bulk.BulkResponse;
@@ -22,13 +23,16 @@
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.action.support.SubscribableListener;
+import org.elasticsearch.action.support.master.AcknowledgedResponse;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.settings.MockSecureSettings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.CollectionUtils;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.ingest.EnterpriseGeoIpTask;
 import org.elasticsearch.ingest.geoip.direct.DatabaseConfiguration;
+import org.elasticsearch.ingest.geoip.direct.DeleteDatabaseConfigurationAction;
 import org.elasticsearch.ingest.geoip.direct.PutDatabaseConfigurationAction;
 import org.elasticsearch.persistent.PersistentTasksService;
 import org.elasticsearch.plugins.Plugin;
@@ -38,12 +42,15 @@
 import org.elasticsearch.test.junit.annotations.TestLogging;
 import org.elasticsearch.transport.RemoteTransportException;
 import org.elasticsearch.xcontent.XContentType;
+import org.junit.After;
 import org.junit.ClassRule;
 
 import java.io.IOException;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 
 import static org.elasticsearch.ingest.EnterpriseGeoIpTask.ENTERPRISE_GEOIP_DOWNLOADER;
 import static org.elasticsearch.ingest.geoip.EnterpriseGeoIpDownloaderTaskExecutor.IPINFO_TOKEN_SETTING;
@@ -54,6 +61,8 @@ public class EnterpriseGeoIpDownloaderIT extends ESIntegTestCase {
 
     private static final String MAXMIND_DATABASE_TYPE = "GeoIP2-City";
     private static final String IPINFO_DATABASE_TYPE = "asn";
+    private static final String MAXMIND_CONFIGURATION = "test-1";
+    private static final String IPINFO_CONFIGURATION = "test-2";
 
     @ClassRule
     public static final EnterpriseGeoIpHttpFixture fixture = new EnterpriseGeoIpHttpFixture(
@@ -118,6 +127,7 @@ public void testEnterpriseDownloaderTask() throws Exception {
          * We know that the databases index has been populated (because we waited around, :wink:), but we don't know for sure that
          * the databases have been pulled down and made available on all nodes. So we run these ingest-and-check steps in assertBusy blocks.
          */
+
         assertBusy(() -> {
             logger.info("Ingesting a test document");
             String documentId = ingestDocument(indexName, geoipPipelineName, sourceField, "89.160.20.128");
@@ -140,6 +150,41 @@ public void testEnterpriseDownloaderTask() throws Exception {
         });
     }
 
+    @After
+    public void cleanup() throws InterruptedException {
+        /*
+         * This method cleans up the database configurations that the test created. This allows the test to be run repeatedly.
+         */
+        CountDownLatch latch = new CountDownLatch(1);
+        LatchedActionListener<AcknowledgedResponse> listener = new LatchedActionListener<>(ActionListener.noop(), latch);
+        SubscribableListener.<AcknowledgedResponse>newForked(
+            l -> admin().cluster()
+                .execute(
+                    DeleteDatabaseConfigurationAction.INSTANCE,
+                    new DeleteDatabaseConfigurationAction.Request(
+                        TimeValue.MAX_VALUE,
+                        TimeValue.timeValueSeconds(10),
+                        MAXMIND_CONFIGURATION
+                    ),
+                    l
+                )
+        )
+            .<AcknowledgedResponse>andThen(
+                l -> admin().cluster()
+                    .execute(
+                        DeleteDatabaseConfigurationAction.INSTANCE,
+                        new DeleteDatabaseConfigurationAction.Request(
+                            TimeValue.MAX_VALUE,
+                            TimeValue.timeValueSeconds(10),
+                            IPINFO_CONFIGURATION
+                        ),
+                        l
+                    )
+            )
+            .addListener(listener);
+        latch.await(10, TimeUnit.SECONDS);
+    }
+
     private void startEnterpriseGeoIpDownloaderTask() {
         PersistentTasksService persistentTasksService = internalCluster().getInstance(PersistentTasksService.class);
         persistentTasksService.sendStartRequest(
@@ -163,7 +208,7 @@ private void configureMaxmindDatabase(String databaseType) {
                 new PutDatabaseConfigurationAction.Request(
                     TimeValue.MAX_VALUE,
                     TimeValue.MAX_VALUE,
-                    new DatabaseConfiguration("test-1", databaseType, new DatabaseConfiguration.Maxmind("test_account"))
+                    new DatabaseConfiguration(MAXMIND_CONFIGURATION, databaseType, new DatabaseConfiguration.Maxmind("test_account"))
                 )
             )
             .actionGet();
@@ -176,7 +221,7 @@ private void configureIpinfoDatabase(String databaseType) {
                 new PutDatabaseConfigurationAction.Request(
                     TimeValue.MAX_VALUE,
                     TimeValue.MAX_VALUE,
-                    new DatabaseConfiguration("test-2", databaseType, new DatabaseConfiguration.Ipinfo())
+                    new DatabaseConfiguration(IPINFO_CONFIGURATION, databaseType, new DatabaseConfiguration.Ipinfo())
                 )
             )
             .actionGet();