Adding a cleanup method to EnterpriseGeoIpDownloaderIT (elastic#125958)

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(
@@ -140,6 +149,28 @@ 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 -> deleteDatabaseConfiguration(MAXMIND_CONFIGURATION, l))
+            .<AcknowledgedResponse>andThen(l -> deleteDatabaseConfiguration(IPINFO_CONFIGURATION, l))
+            .addListener(listener);
+        latch.await(10, TimeUnit.SECONDS);
+    }
+
+    private void deleteDatabaseConfiguration(String configurationName, ActionListener<AcknowledgedResponse> listener) {
+        admin().cluster()
+            .execute(
+                DeleteDatabaseConfigurationAction.INSTANCE,
+                new DeleteDatabaseConfigurationAction.Request(TimeValue.MAX_VALUE, TimeValue.timeValueSeconds(10), configurationName),
+                listener
+            );
+    }
+
     private void startEnterpriseGeoIpDownloaderTask() {
         PersistentTasksService persistentTasksService = internalCluster().getInstance(PersistentTasksService.class);
         persistentTasksService.sendStartRequest(
@@ -163,7 +194,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 +207,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();