Move cancel waiting logic to test functions for DBMAsyncJob

datadog_checks_base/datadog_checks/base/utils/db/utils.py

@@ -256,18 +256,10 @@ def __init__(
 
     def cancel(self):
         """
-        Returns only when the async job loop is fully cancelled.
-        future.result() waits until the timeout (10 seconds) to see that the job loop has exited,
-        so we use it here to verify the async job is cancelled before exiting.
-
-        Raises:
-            TimeoutError if the job loop is not cancelled within 10 seconds
+        Send a signal to cancel the job loop asynchronously.
         """
         self._cancel_event.set()
 
-        if self._job_loop_future is not None:
-            self._job_loop_future.result(10)
-
     def run_job_loop(self, tags):
         """
         :param tags:

datadog_checks_base/tests/base/utils/db/test_util.py

@@ -206,6 +206,7 @@ def test_dbm_async_job_enabled(enabled):
     if enabled:
         assert job._job_loop_future is not None
         job.cancel()
+        job._job_loop_future.result()
     else:
         assert job._job_loop_future is None
 
@@ -215,6 +216,7 @@ def test_dbm_async_job_cancel(aggregator):
     tags = ["hello:there"]
     job.run_job_loop(tags)
     job.cancel()
+    job._job_loop_future.result()
     assert not job._job_loop_future.running(), "thread should be stopped"
     # if the thread doesn't start until after the cancel signal is set then the db connection will never
     # be created in the first place

mysql/tests/test_query_activity.py

@@ -383,6 +383,7 @@ def test_async_job_enabled(dd_run_check, dbm_instance, activity_enabled):
     check.cancel()
     if activity_enabled:
         assert check._query_activity._job_loop_future is not None
+        check._query_activity._job_loop_future.result()
     else:
         assert check._query_activity._job_loop_future is None
 
@@ -408,6 +409,8 @@ def test_async_job_cancel(aggregator, dd_run_check, dbm_instance):
     check = MySql(CHECK_NAME, {}, [dbm_instance])
     dd_run_check(check)
     check.cancel()
+    # wait for it to stop and make sure it doesn't throw any exceptions
+    check._query_activity._job_loop_future.result()
     assert not check._query_activity._job_loop_future.running(), "activity thread should be stopped"
     # if the thread doesn't start until after the cancel signal is set then the db connection will never
     # be created in the first place

mysql/tests/test_statements.py

@@ -809,6 +809,9 @@ def test_async_job_cancel(aggregator, dd_run_check, dbm_instance):
     mysql_check = MySql(common.CHECK_NAME, {}, [dbm_instance])
     dd_run_check(mysql_check)
     mysql_check.cancel()
+    # wait for it to stop and make sure it doesn't throw any exceptions
+    mysql_check._statement_samples._job_loop_future.result()
+    mysql_check._statement_metrics._job_loop_future.result()
     assert not mysql_check._statement_samples._job_loop_future.running(), "samples thread should be stopped"
     assert not mysql_check._statement_metrics._job_loop_future.running(), "metrics thread should be stopped"
     assert mysql_check._statement_samples._db is None, "samples db connection should be gone"
@@ -844,10 +847,12 @@ def test_async_job_enabled(dd_run_check, dbm_instance, statement_samples_enabled
     mysql_check.cancel()
     if statement_samples_enabled:
         assert mysql_check._statement_samples._job_loop_future is not None
+        mysql_check._statement_samples._job_loop_future.result()
     else:
         assert mysql_check._statement_samples._job_loop_future is None
     if statement_metrics_enabled:
         assert mysql_check._statement_metrics._job_loop_future is not None
+        mysql_check._statement_metrics._job_loop_future.result()
     else:
         assert mysql_check._statement_metrics._job_loop_future is None
 

postgres/tests/utils.py

@@ -54,3 +54,9 @@ def run_one_check(check, db_instance):
     """
     check.check(db_instance)
     check.cancel()
+    if check.statement_samples._job_loop_future is not None:
+        check.statement_samples._job_loop_future.result()
+    if check.statement_metrics._job_loop_future is not None:
+        check.statement_metrics._job_loop_future.result()
+    if check.metadata_samples._job_loop_future is not None:
+        check.metadata_samples._job_loop_future.result()

sqlserver/tests/test_activity.py

@@ -716,6 +716,7 @@ def test_async_job_enabled(dd_run_check, dbm_instance, activity_enabled):
     check.cancel()
     if activity_enabled:
         assert check.activity._job_loop_future is not None
+        check.activity._job_loop_future.result()
     else:
         assert check.activity._job_loop_future is None
 
@@ -741,6 +742,8 @@ def test_async_job_cancel_cancel(aggregator, dd_run_check, dbm_instance):
     check = SQLServer(CHECK_NAME, {}, [dbm_instance])
     dd_run_check(check)
     check.cancel()
+    # wait for it to stop and make sure it doesn't throw any exceptions
+    check.activity._job_loop_future.result()
     assert not check.activity._job_loop_future.running(), "activity thread should be stopped"
     # if the thread doesn't start until after the cancel signal is set then the db connection will never
     # be created in the first place

sqlserver/tests/test_statements.py

@@ -773,6 +773,7 @@ def test_async_job_enabled(dd_run_check, dbm_instance, statement_metrics_enabled
     check.cancel()
     if statement_metrics_enabled:
         assert check.statement_metrics._job_loop_future is not None
+        check.statement_metrics._job_loop_future.result()
     else:
         assert check.statement_metrics._job_loop_future is None
 
@@ -798,6 +799,8 @@ def test_async_job_cancel_cancel(aggregator, dd_run_check, dbm_instance):
     check = SQLServer(CHECK_NAME, {}, [dbm_instance])
     dd_run_check(check)
     check.cancel()
+    # wait for it to stop and make sure it doesn't throw any exceptions
+    check.statement_metrics._job_loop_future.result()
     assert not check.statement_metrics._job_loop_future.running(), "metrics thread should be stopped"
     # if the thread doesn't start until after the cancel signal is set then the db connection will never
     # be created in the first place