From ae64fed56a07e57b1f5330d14105e77ace55fe4c Mon Sep 17 00:00:00 2001
From: Codegass <wchweichenhao@gmail.com>
Date: Mon, 13 May 2024 17:04:58 -0400
Subject: [PATCH] Add Assertion to the `SharedTimerTest` (#2381)

---
 .../sqlserver/jdbc/SharedTimerTest.java       | 20 ++++++++++++++++---
 .../sqlserver/jdbc/TestResource.java          |  1 +
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/SharedTimerTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/SharedTimerTest.java
index 6aa8bff2e..48a70dd29 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/SharedTimerTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/SharedTimerTest.java
@@ -1,9 +1,16 @@
 package com.microsoft.sqlserver.jdbc;
 
-import org.junit.jupiter.api.Test;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 
 import java.util.ArrayList;
-import java.util.concurrent.*;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import org.junit.jupiter.api.Test;
 
 
 class SharedTimerTest {
@@ -11,8 +18,8 @@ class SharedTimerTest {
     @Test
     void getTimer() throws InterruptedException, ExecutionException, TimeoutException {
         final int iterations = 500;
-
         ExecutorService executor = Executors.newFixedThreadPool(2);
+
         try {
             ArrayList<CompletableFuture<?>> futures = new ArrayList<>(iterations);
             for (int i = 0; i < iterations; i++) {
@@ -22,6 +29,13 @@ void getTimer() throws InterruptedException, ExecutionException, TimeoutExceptio
             CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).get(2, TimeUnit.MINUTES);
         } finally {
             executor.shutdown();
+            // 5000ms wait time for the AzureDB connection to close, need full test in the
+            // test lab for the exact time
+            if (!executor.awaitTermination(5000, TimeUnit.MILLISECONDS)) {
+                executor.shutdownNow();
+            }
         }
+
+        assertFalse(SharedTimer.isRunning(), TestResource.getResource("R_sharedTimerStopOnNoRef"));
     }
 }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java b/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java
index d542434fb..a630f5e0d 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java
@@ -212,6 +212,7 @@ protected Object[][] getContents() {
             {"R_objectNullOrEmpty", "The {0} is null or empty."},
             {"R_cekDecryptionFailed", "Failed to decrypt a column encryption key using key store provider: {0}."},
             {"R_connectTimedOut", "connect timed out"},
+            {"R_sharedTimerStopOnNoRef", "SharedTimer should be stopped after all references are removed."},
             {"R_sessionKilled", "Cannot continue the execution because the session is in the kill state"},
             {"R_failedFedauth", "Failed to acquire fedauth token: "},
             {"R_noLoginModulesConfiguredForJdbcDriver",