Make sure async_policy_saver gets closed in tests to avoid hanging in…

… OSS tests.

PiperOrigin-RevId: 302460421
Change-Id: Ie1a52a7b29b9bb25da7dec69ad1a89273f2392f0

broken_tests.txt

@@ -1,2 +1 @@
 replay_buffers.tfrecord_replay_buffer_test # b/140896267
-policies.async_policy_saver_test # b/151865318

tf_agents/policies/async_policy_saver_test.py

@@ -38,6 +38,9 @@ def testSave(self):
     async_saver.flush()
 
     saver.save.assert_called_once_with(save_path)
+    # Have to close the saver to avoid hanging threads that will prevent OSS
+    # tests from finishing.
+    async_saver.close()
 
   def testCheckpointSave(self):
     saver = mock.create_autospec(policy_saver.PolicySaver, instance=True)
@@ -52,6 +55,9 @@ def testCheckpointSave(self):
     async_saver.flush()
 
     saver.save_checkpoint.assert_called_once_with(checkpoint_path)
+    # Have to close the saver to avoid hanging threads that will prevent OSS
+    # tests from finishing.
+    async_saver.close()
 
   def testBlockingSave(self):
     saver = mock.create_autospec(policy_saver.PolicySaver, instance=True)
@@ -64,6 +70,9 @@ def testBlockingSave(self):
     async_saver.save(path2, blocking=True)
 
     saver.save.assert_has_calls([mock.call(path1), mock.call(path2)])
+    # Have to close the saver to avoid hanging threads that will prevent OSS
+    # tests from finishing.
+    async_saver.close()
 
   def testBlockingCheckpointSave(self):
     saver = mock.create_autospec(policy_saver.PolicySaver, instance=True)
@@ -76,6 +85,9 @@ def testBlockingCheckpointSave(self):
     async_saver.save_checkpoint(path2, blocking=True)
 
     saver.save_checkpoint.assert_has_calls([mock.call(path1), mock.call(path2)])
+    # Have to close the saver to avoid hanging threads that will prevent OSS
+    # tests from finishing.
+    async_saver.close()
 
   def testClose(self):
     saver = mock.create_autospec(policy_saver.PolicySaver, instance=True)