From 48e65c486205e0ebbda485652347af42ac202f34 Mon Sep 17 00:00:00 2001
From: Matthew Elwell <matthew.elwell@flagsmith.com>
Date: Tue, 5 Dec 2023 09:42:13 +0000
Subject: [PATCH] fix: FlagsmithClient.close() doesn't kill polling manager
 properly (#133)

* Convert to daemon thread and ensure that closing works correctly

* Remove unnecessary call to getEnvironmentFlags

* Add comment

* typo

Co-authored-by: Kim Gustyr <kim.gustyr@flagsmith.com>

---------

Co-authored-by: Kim Gustyr <kim.gustyr@flagsmith.com>
---
 .../com/flagsmith/threads/PollingManager.java | 11 ++++--
 .../com/flagsmith/FlagsmithClientTest.java    | 36 ++++++++++++++++++-
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/src/main/java/com/flagsmith/threads/PollingManager.java b/src/main/java/com/flagsmith/threads/PollingManager.java
index 56637c1e..76090576 100644
--- a/src/main/java/com/flagsmith/threads/PollingManager.java
+++ b/src/main/java/com/flagsmith/threads/PollingManager.java
@@ -32,7 +32,7 @@ public PollingManager(FlagsmithClient client, Integer interval) {
    * @return
    */
   private Thread initializeThread() {
-    return new Thread() {
+    Thread thread = new Thread() {
       @Override
       public void run() {
         try {
@@ -45,6 +45,8 @@ public void run() {
         }
       }
     };
+    thread.setDaemon(true);
+    return thread;
   }
 
   /**
@@ -61,5 +63,10 @@ public void stopPolling() {
     internalThread.interrupt();
   }
 
-
+  /**
+   * Get thread status
+   */
+  public Boolean getIsThreadAlive() {
+    return internalThread.isAlive();
+  }
 }
diff --git a/src/test/java/com/flagsmith/FlagsmithClientTest.java b/src/test/java/com/flagsmith/FlagsmithClientTest.java
index 4059a399..fc8d5a7a 100644
--- a/src/test/java/com/flagsmith/FlagsmithClientTest.java
+++ b/src/test/java/com/flagsmith/FlagsmithClientTest.java
@@ -7,11 +7,12 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.JsonNode;
-import com.flagsmith.responses.FlagsAndTraitsResponse;
 import com.flagsmith.config.FlagsmithCacheConfig;
 import com.flagsmith.config.FlagsmithConfig;
 import com.flagsmith.exceptions.FlagsmithApiError;
@@ -24,6 +25,7 @@
 import com.flagsmith.models.DefaultFlag;
 import com.flagsmith.models.Flags;
 import com.flagsmith.models.Segment;
+import com.flagsmith.responses.FlagsAndTraitsResponse;
 import com.flagsmith.threads.PollingManager;
 import com.flagsmith.threads.RequestProcessor;
 
@@ -692,4 +694,36 @@ public void testGetIdentityFlags_UsesDefaultFlags_IfLocalEvaluationEnvironmentNu
         assertEquals(identityFlags.getFeatureValue("foo"), DEFAULT_FLAG_VALUE);
         assertEquals(identityFlags.isFeatureEnabled("foo"), DEFAULT_FLAG_STATE);
     }
+
+    @Test
+    public void testClose() throws FlagsmithApiError, InterruptedException {
+        // Given
+        int pollingInterval = 1;
+
+        FlagsmithConfig config = FlagsmithConfig
+                .newBuilder()
+                .withLocalEvaluation(true)
+                .withEnvironmentRefreshIntervalSeconds(pollingInterval)
+                .build();
+
+        FlagsmithApiWrapper mockedApiWrapper = mock(FlagsmithApiWrapper.class);
+        when(mockedApiWrapper.getEnvironment()).thenReturn(FlagsmithTestHelper.environmentModel());
+        when(mockedApiWrapper.getConfig()).thenReturn(config);
+
+        FlagsmithClient client = FlagsmithClient.newBuilder()
+                .withFlagsmithApiWrapper(mockedApiWrapper)
+                .withConfiguration(config)
+                .setApiKey("ser.dummy-key")
+                .build();
+
+        // When
+        client.close();
+
+        // Then
+        // Since the thread will only stop once it reads the interrupt signal correctly
+        // on its next polling interval, we need to wait for the polling interval
+        // to complete before checking the thread has been killed correctly.
+        Thread.sleep(pollingInterval);
+        assertFalse(client.getPollingManager().getIsThreadAlive());
+    }
 }