From e9e7d56df3042b02882000dc0a7f045c01194400 Mon Sep 17 00:00:00 2001
From: Szymon Bialkowski <szybia@tuta.io>
Date: Thu, 21 Aug 2025 16:07:33 +0200
Subject: [PATCH] Re-enable EnterpriseGeoIpDownloaderIT test

---
 .../geoip/EnterpriseGeoIpDownloaderIT.java    | 35 +++++++++++++++----
 muted-tests.yml                               |  3 --
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/modules/ingest-geoip/src/internalClusterTest/java/org/elasticsearch/ingest/geoip/EnterpriseGeoIpDownloaderIT.java b/modules/ingest-geoip/src/internalClusterTest/java/org/elasticsearch/ingest/geoip/EnterpriseGeoIpDownloaderIT.java
index 4fdee727b5755..8afc2ddd8fcf0 100644
--- a/modules/ingest-geoip/src/internalClusterTest/java/org/elasticsearch/ingest/geoip/EnterpriseGeoIpDownloaderIT.java
+++ b/modules/ingest-geoip/src/internalClusterTest/java/org/elasticsearch/ingest/geoip/EnterpriseGeoIpDownloaderIT.java
@@ -47,11 +47,13 @@
 import org.junit.ClassRule;
 
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.elasticsearch.ingest.EnterpriseGeoIpTask.ENTERPRISE_GEOIP_DOWNLOADER;
 import static org.elasticsearch.ingest.geoip.EnterpriseGeoIpDownloaderTaskExecutor.IPINFO_TOKEN_SETTING;
@@ -151,7 +153,7 @@ public void testEnterpriseDownloaderTask() throws Exception {
     }
 
     @After
-    public void cleanup() throws InterruptedException {
+    public void cleanup() throws Exception {
         /*
          * This method cleans up the database configurations that the test created. This allows the test to be run repeatedly.
          */
@@ -161,6 +163,18 @@ public void cleanup() throws InterruptedException {
             .<AcknowledgedResponse>andThen(l -> deleteDatabaseConfiguration(IPINFO_CONFIGURATION, l))
             .addListener(listener);
         latch.await(10, TimeUnit.SECONDS);
+        PersistentTasksService persistentTasksService = internalCluster().getInstance(PersistentTasksService.class);
+        AtomicBoolean removeProjectSucceeded = new AtomicBoolean(false);
+        persistentTasksService.sendProjectRemoveRequest(
+            ProjectId.DEFAULT,
+            ENTERPRISE_GEOIP_DOWNLOADER,
+            TimeValue.MAX_VALUE,
+            ActionListener.wrap(r -> {
+                logger.info("stopped enterprise geoip downloader task");
+                removeProjectSucceeded.set(true);
+            }, e -> { fail(e, "Failed to stop enterprise geoip downloader task"); })
+        );
+        assertBusy(() -> assertTrue(removeProjectSucceeded.get()));
     }
 
     private void deleteDatabaseConfiguration(String configurationName, ActionListener<AcknowledgedResponse> listener) {
@@ -172,21 +186,29 @@ private void deleteDatabaseConfiguration(String configurationName, ActionListene
             );
     }
 
-    private void startEnterpriseGeoIpDownloaderTask(ProjectId projectId) {
+    private void startEnterpriseGeoIpDownloaderTask(ProjectId projectId) throws Exception {
         PersistentTasksService persistentTasksService = internalCluster().getInstance(PersistentTasksService.class);
+        AtomicBoolean completed = new AtomicBoolean(false);
         persistentTasksService.sendProjectStartRequest(
             projectId,
             ENTERPRISE_GEOIP_DOWNLOADER,
             ENTERPRISE_GEOIP_DOWNLOADER,
             new EnterpriseGeoIpTask.EnterpriseGeoIpTaskParams(),
             TimeValue.MAX_VALUE,
-            ActionListener.wrap(r -> logger.debug("Started enterprise geoip downloader task"), e -> {
+            ActionListener.wrap(r -> {
+                logger.info("Started enterprise geoip downloader task");
+                completed.set(true);
+            }, e -> {
                 Throwable t = e instanceof RemoteTransportException ? ExceptionsHelper.unwrapCause(e) : e;
-                if (t instanceof ResourceAlreadyExistsException == false) {
-                    logger.error("failed to create enterprise geoip downloader task", e);
+                if (t instanceof ResourceAlreadyExistsException) {
+                    logger.info("enterprise geoip downloader task already started");
+                    completed.set(true);
+                } else {
+                    fail(e, "Failed to create enterprise geoip downloader task");
                 }
             })
         );
+        assertBusy(() -> assertTrue(completed.get()));
     }
 
     private void configureMaxmindDatabase(String databaseType) {
@@ -220,7 +242,8 @@ private void waitAround() throws Exception {
         assertBusy(() -> {
             SearchResponse searchResponse = client().search(new SearchRequest(GeoIpDownloader.DATABASES_INDEX)).actionGet();
             try {
-                assertThat(searchResponse.getHits().getHits().length, equalTo(2));
+                final var hits = searchResponse.getHits().getHits();
+                assertThat(Arrays.toString(hits), hits.length, equalTo(2));
             } finally {
                 searchResponse.decRef();
             }
diff --git a/muted-tests.yml b/muted-tests.yml
index 4ef53e040944b..eb91d837fd642 100644
--- a/muted-tests.yml
+++ b/muted-tests.yml
@@ -230,9 +230,6 @@ tests:
 - class: org.elasticsearch.smoketest.MlWithSecurityIT
   method: test {yaml=ml/trained_model_cat_apis/Test cat trained models}
   issue: https://github.com/elastic/elasticsearch/issues/125750
-- class: org.elasticsearch.ingest.geoip.EnterpriseGeoIpDownloaderIT
-  method: testEnterpriseDownloaderTask
-  issue: https://github.com/elastic/elasticsearch/issues/126124
 - class: org.elasticsearch.xpack.test.rest.XPackRestIT
   method: test {p0=transform/transforms_start_stop/Test start/stop only starts/stops specified transform}
   issue: https://github.com/elastic/elasticsearch/issues/126466