fix mypy errors

Author: tmonty12

tests/fault_tolerance/deploy/client.py

@@ -25,6 +25,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import requests
+from kr8s.objects import Pod
 
 from tests.utils.managed_deployment import ManagedDeployment
 
@@ -45,7 +46,7 @@ def get_frontend_port(
     deployment_spec: Any,
     pod_ports: Dict[str, Any],
     logger: logging.Logger,
-) -> Tuple[Optional[str], Optional[int], Optional[str]]:
+) -> Tuple[Optional[str], Optional[int], Optional[Pod]]:
     """
     Select a frontend pod using round-robin and setup port forwarding.
 

tests/fault_tolerance/deploy/scenarios.py

@@ -199,7 +199,7 @@ class RollingUpgradeFailure(Failure):
 
     async def execute(
         self, deployment: ManagedDeployment, logger: logging.Logger
-    ) -> None:
+    ) -> list[str]:
         """Execute rolling upgrade failure injection."""
         await deployment.trigger_rolling_upgrade(self.service_names)
 
@@ -225,7 +225,7 @@ class DeletePodFailure(Failure):
 
     async def execute(
         self, deployment: ManagedDeployment, logger: logging.Logger
-    ) -> None:
+    ) -> list[str]:
         """Execute pod deletion failure injection."""
         service_pod_dict = deployment.get_pods(self.service_names)
         pod_names: list[str] = []
@@ -276,7 +276,7 @@ def __init__(
 
     async def execute(
         self, deployment: ManagedDeployment, logger: logging.Logger
-    ) -> None:
+    ) -> list[str]:
         """Execute process termination failure injection."""
         service_pod_dict = deployment.get_pods(self.service_names)
         pod_names: list[str] = []
@@ -323,7 +323,7 @@ def __init__(
 
     async def execute(
         self, deployment: ManagedDeployment, logger: logging.Logger
-    ) -> None:
+    ) -> list[str]:
         """Token overflow is handled client-side, so this is a no-op."""
         # The actual overflow is handled by the client configuration
         # which uses the input_token_length from the Load config

tests/fault_tolerance/deploy/test_deployment.py

@@ -9,7 +9,7 @@
 import signal
 from contextlib import contextmanager
 from multiprocessing.context import SpawnProcess
-from typing import Any
+from typing import Any, Optional
 
 import pytest
 
@@ -463,6 +463,7 @@ async def test_fault_scenario(
     if image:
         scenario.deployment.set_image(image)
 
+    model: Optional[str] = None
     if scenario.model:
         scenario.deployment.set_model(scenario.model)
         model = scenario.model

tests/utils/managed_deployment.py

@@ -57,6 +57,14 @@ def image(self) -> Optional[str]:
         except KeyError:
             return None
 
+    @image.setter
+    def image(self, value: str):
+        if "extraPodSpec" not in self._spec:
+            self._spec["extraPodSpec"] = {"mainContainer": {}}
+        if "mainContainer" not in self._spec["extraPodSpec"]:
+            self._spec["extraPodSpec"]["mainContainer"] = {}
+        self._spec["extraPodSpec"]["mainContainer"]["image"] = value
+
     @property
     def envs(self) -> list[dict[str, str]]:
         """Environment variables for the service"""
@@ -66,14 +74,6 @@ def envs(self) -> list[dict[str, str]]:
     def envs(self, value: list[dict[str, str]]):
         self._spec["envs"] = value
 
-    @image.setter
-    def image(self, value: str):
-        if "extraPodSpec" not in self._spec:
-            self._spec["extraPodSpec"] = {"mainContainer": {}}
-        if "mainContainer" not in self._spec["extraPodSpec"]:
-            self._spec["extraPodSpec"]["mainContainer"] = {}
-        self._spec["extraPodSpec"]["mainContainer"]["image"] = value
-
     # ----- Replicas -----
     @property
     def replicas(self) -> int:
@@ -729,7 +729,7 @@ async def trigger_rolling_upgrade(self, service_names: list[str]):
                 "service_names cannot be empty for trigger_rolling_upgrade"
             )
 
-        patch_body = {"spec": {"services": {}}}
+        patch_body: dict[str, Any] = {"spec": {"services": {}}}
 
         for service_name in service_names:
             self.deployment_spec.set_service_env_var(