Splitting time consuming tests to multiple ones

Creating and closing processes and threads can be quite time consuming resulting to test timeouts if the tests performs a lot of actions.
ilastik · Dec 9, 2024 · cbb2619 · cbb2619
1 parent 6d6749f
commit cbb2619
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 8 deletions.
diff --git a/pytest.ini b/pytest.ini
@@ -1,7 +1,7 @@
 [pytest]
 python_files = test_*.py
 addopts =
-    --timeout 10
+    --timeout 15
     -v
     -s
     --color=yes

diff --git a/tests/test_server/test_grpc/test_training_servicer.py b/tests/test_server/test_grpc/test_training_servicer.py
@@ -403,34 +403,39 @@ def fit(self):
         pause_thread.join()
         resume_thread.join()
 
-    def test_graceful_shutdown_for_any_state(self, grpc_stub):
-        # after init
+    def test_graceful_shutdown_after_init(self, grpc_stub):
         init_response = grpc_stub.Init(training_pb2.TrainingConfig(yaml_content=prepare_unet2d_test_environment()))
         training_session_id = training_pb2.TrainingSessionId(id=init_response.id)
         grpc_stub.CloseTrainerSession(training_session_id)
 
-        # after start
+    def test_graceful_shutdown_after_start(self, grpc_stub):
         init_response = grpc_stub.Init(training_pb2.TrainingConfig(yaml_content=prepare_unet2d_test_environment()))
         training_session_id = training_pb2.TrainingSessionId(id=init_response.id)
         grpc_stub.Start(training_session_id)
         grpc_stub.CloseTrainerSession(training_session_id)
 
-        # after pause
+    def test_graceful_shutdown_after_pause(self, grpc_stub):
         init_response = grpc_stub.Init(training_pb2.TrainingConfig(yaml_content=prepare_unet2d_test_environment()))
         training_session_id = training_pb2.TrainingSessionId(id=init_response.id)
         grpc_stub.Start(training_session_id)
         grpc_stub.Pause(training_session_id)
         grpc_stub.CloseTrainerSession(training_session_id)
 
-        # after resume
+    def test_graceful_shutdown_after_resume(self, grpc_stub):
         init_response = grpc_stub.Init(training_pb2.TrainingConfig(yaml_content=prepare_unet2d_test_environment()))
         training_session_id = training_pb2.TrainingSessionId(id=init_response.id)
         grpc_stub.Start(training_session_id)
         grpc_stub.Pause(training_session_id)
         grpc_stub.Resume(training_session_id)
         grpc_stub.CloseTrainerSession(training_session_id)
 
-        # attempt to close again
+    def test_close_trainer_session_twice(self, grpc_stub):
+        # Attempt to close the session twice
+        init_response = grpc_stub.Init(training_pb2.TrainingConfig(yaml_content=prepare_unet2d_test_environment()))
+        training_session_id = training_pb2.TrainingSessionId(id=init_response.id)
+        grpc_stub.CloseTrainerSession(training_session_id)
+
+        # The second attempt should raise an error
         with pytest.raises(grpc.RpcError) as excinfo:
             grpc_stub.CloseTrainerSession(training_session_id)
         assert "Unknown session" in excinfo.value.details()

diff --git a/tiktorch/server/session/backend/supervisor.py b/tiktorch/server/session/backend/supervisor.py
@@ -46,7 +46,6 @@ def __init__(self, trainer: Trainer) -> None:
         super().__init__()
         self._trainer = trainer
         self._trainer.should_stop_callbacks.register(self._should_stop)
-        self._shutdown_event = threading.Event()
         self._state = TrainerState.IDLE
         self._pause_triggered = False
         self._session_thread = threading.Thread(target=self._start_session, name="SessionThread")