Type fixes

Author: jakozaur

src/cli/pentest.py

@@ -213,10 +213,14 @@ def main(
     click.echo(f"✅ {backend_type} model {client.get_model_name()} ready")
 
     # Check if backend is busy before starting tests (Ollama only)
-    if not skip_busy_check and hasattr(client, "check_status"):
+    if (
+        not skip_busy_check
+        and hasattr(client, "check_status")
+        and callable(getattr(client, "check_status", None))
+    ):
         click.echo(f"🔍 Checking {backend_type} status...")
         try:
-            status = client.check_status()
+            status = client.check_status()  # type: ignore
 
             if status.is_busy:
                 click.echo(f"⚠️  WARNING: {backend_type} appears busy!")

src/utils/model_client.py

@@ -2,17 +2,21 @@
 
 import subprocess
 import time
-from typing import Any
+from typing import TYPE_CHECKING, Any, Union
 
 import requests
 from requests.exceptions import Timeout
 from src.models import ModelResponse, OllamaStatus
-from src.utils.llm_backend import create_backend
 from src.utils.settings_manager import settings_manager
 
+if TYPE_CHECKING:
+    from src.utils.llm_backend import LLMBackend
 
-def get_client(seed: int | None = None) -> object:
+
+def get_client(seed: int | None = None) -> Union["LLMBackend", "OllamaClient"]:
     """Factory function to get the configured LLM client."""
+    from src.utils.llm_backend import create_backend
+
     try:
         settings = settings_manager.load_settings()
         return create_backend(settings, seed)
@@ -313,6 +317,18 @@ def get_backend_type(self) -> str:
         """Get the backend type identifier (for compatibility)."""
         return "Ollama"
 
+    def get_model_name(self) -> str:
+        """Get the model name (for compatibility)."""
+        return self.model
+
+    def is_available(self) -> bool:
+        """Check if model is available (for compatibility)."""
+        return self.is_model_available()
+
+    def check_status(self) -> OllamaStatus:
+        """Check Ollama status (for compatibility)."""
+        return self.check_ollama_status()
+
 
 def test_connection() -> bool | None:
     """Test Ollama connection and model availability"""

tests/test_model_client.py

@@ -21,56 +21,54 @@ def test_default_initialization(self) -> None:
         client = OllamaClient()
         assert client.base_url == "http://localhost:11434"
         assert client.model == "gpt-oss:20b"
-        assert client.session is not None
 
     def test_custom_initialization(self) -> None:
         """Test client accepts custom configuration"""
         client = OllamaClient(host="custom.host", port=8080, model="custom:model")
         assert client.base_url == "http://custom.host:8080"
         assert client.model == "custom:model"
 
-    def test_session_persistence(self) -> None:
-        """Test that session object persists across calls"""
-        client = OllamaClient()
-        session1 = client.session
-        session2 = client.session
-        assert session1 is session2
+    def test_seed_initialization(self) -> None:
+        """Test that seed parameter is properly stored"""
+        client = OllamaClient(seed=42)
+        assert client.seed == 42
+        
+        client_no_seed = OllamaClient()
+        assert client_no_seed.seed is None
 
 
 class TestOllamaClientModelAvailability:
     """Test model availability checking"""
 
-    @patch('src.utils.model_client.requests.Session')
-    def test_model_available(self, mock_session_class) -> None:
+    @patch('src.utils.model_client.requests.get')
+    def test_model_available(self, mock_get) -> None:
         """Test when model is available"""
-        mock_session = MagicMock()
         mock_response = MagicMock()
         mock_response.json.return_value = {
             "models": [
                 {"name": "gpt-oss:20b"},
                 {"name": "other:model"}
             ]
         }
-        mock_session.get.return_value = mock_response
-        mock_session_class.return_value = mock_session
+        mock_response.raise_for_status.return_value = None
+        mock_get.return_value = mock_response
 
         client = OllamaClient()
         assert client.is_model_available() is True
-        mock_session.get.assert_called_once_with(
+        mock_get.assert_called_once_with(
             "http://localhost:11434/api/tags",
             timeout=180
         )
 
-    @patch('src.utils.model_client.requests.Session')
-    def test_model_not_available(self, mock_session_class) -> None:
+    @patch('src.utils.model_client.requests.get')
+    def test_model_not_available(self, mock_get) -> None:
         """Test when model is not available"""
-        mock_session = MagicMock()
         mock_response = MagicMock()
         mock_response.json.return_value = {
             "models": [{"name": "other:model"}]
         }
-        mock_session.get.return_value = mock_response
-        mock_session_class.return_value = mock_session
+        mock_response.raise_for_status.return_value = None
+        mock_get.return_value = mock_response
 
         client = OllamaClient()
         assert client.is_model_available() is False