Mypy and unit test fixes

multiscale · Aug 13, 2024 · 523d1fd · 523d1fd
1 parent 303e2d9
commit 523d1fd
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 10 deletions.
diff --git a/libmuscle/python/libmuscle/manager/deadlock_detector.py b/libmuscle/python/libmuscle/manager/deadlock_detector.py
@@ -21,7 +21,7 @@ def __init__(
 
         self._queue: Queue[Optional[_QueueItem]] = Queue()
         self._waiting_instances: Dict[str, str] = {}
-        self._waiting_instance_ports: Dict[str, Tuple[str, int]] = {}
+        self._waiting_instance_ports: Dict[str, Tuple[str, Optional[int]]] = {}
 
         self._detected_deadlocks: List[str] = []
 

diff --git a/libmuscle/python/libmuscle/manager/test/conftest.py b/libmuscle/python/libmuscle/manager/test/conftest.py
@@ -8,6 +8,7 @@
 from libmuscle.manager.snapshot_registry import SnapshotRegistry
 from libmuscle.manager.topology_store import TopologyStore
 from libmuscle.manager.profile_store import ProfileStore
+from libmuscle.manager.deadlock_detector import DeadlockDetector
 
 
 @pytest.fixture
@@ -54,13 +55,18 @@ def snapshot_registry(mmp_configuration, topology_store) -> SnapshotRegistry:
     return SnapshotRegistry(mmp_configuration, None, topology_store)
 
 
+@pytest.fixture
+def deadlock_detector() -> DeadlockDetector:
+    return DeadlockDetector(lambda: None, 1.0)
+
+
 @pytest.fixture
 def mmp_request_handler(
         logger, profile_store, mmp_configuration, instance_registry,
-        topology_store, snapshot_registry):
+        topology_store, snapshot_registry, deadlock_detector):
     return MMPRequestHandler(
             logger, profile_store, mmp_configuration, instance_registry,
-            topology_store, snapshot_registry, None)
+            topology_store, snapshot_registry, deadlock_detector, None)
 
 
 @pytest.fixture
@@ -77,10 +83,10 @@ def loaded_instance_registry(instance_registry):
 @pytest.fixture
 def registered_mmp_request_handler(
         logger, profile_store, mmp_configuration, loaded_instance_registry,
-        topology_store, snapshot_registry):
+        topology_store, snapshot_registry, deadlock_detector):
     return MMPRequestHandler(
             logger, profile_store, mmp_configuration, loaded_instance_registry,
-            topology_store, snapshot_registry, None)
+            topology_store, snapshot_registry, deadlock_detector, None)
 
 
 @pytest.fixture
@@ -133,8 +139,8 @@ def loaded_instance_registry2():
 @pytest.fixture
 def registered_mmp_request_handler2(
         logger, profile_store, mmp_configuration, loaded_instance_registry2,
-        topology_store2, snapshot_registry2, tmp_path):
+        topology_store2, snapshot_registry2, deadlock_detector, tmp_path):
     return MMPRequestHandler(
             logger, profile_store, mmp_configuration,
             loaded_instance_registry2, topology_store2, snapshot_registry2,
-            RunDir(tmp_path))
+            deadlock_detector, RunDir(tmp_path))
diff --git a/libmuscle/python/libmuscle/manager/test/test_mmp_request_handler.py b/libmuscle/python/libmuscle/manager/test/test_mmp_request_handler.py
@@ -15,10 +15,10 @@
 
 def test_create_servicer(
         logger, profile_store, mmp_configuration, instance_registry,
-        topology_store, snapshot_registry):
+        topology_store, snapshot_registry, deadlock_detector):
     MMPRequestHandler(
             logger, profile_store, mmp_configuration, instance_registry,
-            topology_store, snapshot_registry, None)
+            topology_store, snapshot_registry, deadlock_detector, None)
 
 
 def test_log_message(mmp_request_handler, caplog):

diff --git a/libmuscle/python/libmuscle/mcp/tcp_transport_client.py b/libmuscle/python/libmuscle/mcp/tcp_transport_client.py
@@ -53,7 +53,7 @@ def __init__(self, location: str) -> None:
             self._socket = sock
 
         if hasattr(select, "poll"):
-            self._poll_obj = select.poll()
+            self._poll_obj: Optional[select.poll] = select.poll()
             self._poll_obj.register(self._socket, select.POLLIN)
         else:
             self._poll_obj = None  # On platforms that don't support select.poll
@@ -84,6 +84,7 @@ def call(self, request: bytes, timeout_handler: Optional[TimeoutHandler] = None
 
         length = recv_int64(self._socket)
         if did_timeout:
+            assert timeout_handler is not None  # mypy
             timeout_handler.on_receive()
         start_transfer = ProfileTimestamp()