fix test

2 files changed, 7 insertions, 0 deletions

diff --git a/src/python/grpcio_tests/tests/health_check/BUILD.bazel b/src/python/grpcio_tests/tests/health_check/BUILD.bazel
index 19e1e1b2e1..77bc61aa30 100644
--- a/src/python/grpcio_tests/tests/health_check/BUILD.bazel
+++ b/src/python/grpcio_tests/tests/health_check/BUILD.bazel
@@ -9,6 +9,7 @@ py_test(
         "//src/python/grpcio/grpc:grpcio",
         "//src/python/grpcio_health_checking/grpc_health/v1:grpc_health",
         "//src/python/grpcio_tests/tests/unit:test_common",
+        "//src/python/grpcio_tests/tests/unit/framework/common:common",
     ],
     imports = ["../../",],
 )
diff --git a/src/python/grpcio_tests/tests/health_check/_health_servicer_test.py b/src/python/grpcio_tests/tests/health_check/_health_servicer_test.py
index bf90fa15c0..35794987bc 100644
--- a/src/python/grpcio_tests/tests/health_check/_health_servicer_test.py
+++ b/src/python/grpcio_tests/tests/health_check/_health_servicer_test.py
@@ -14,6 +14,7 @@
 """Tests of grpc_health.v1.health."""
 
 import threading
+import time
 import unittest
 
 import grpc
@@ -206,6 +207,11 @@ class HealthServicerTest(unittest.TestCase):
         self._servicer.set(_WATCH_SERVICE,
                            health_pb2.HealthCheckResponse.SERVING)
         thread.join()
+
+        # Wait, if necessary, for serving thread to process client cancellation
+        timeout = time.time() + test_constants.SHORT_TIMEOUT
+        while time.time() < timeout and self._servicer._watchers[_WATCH_SERVICE]:
+            time.sleep(1)
         self.assertFalse(self._servicer._watchers[_WATCH_SERVICE],
                          'watch set should be empty')
         self.assertTrue(response_queue.empty())