Fix linter warnings and clean up tests

Author: SaptaDey

pyrightconfig.json

@@ -21,6 +21,5 @@
     "reportUnusedVariable": false,
     "reportDuplicateImport": true,
     "reportGeneralTypeIssues": false,
--    "reportPrivateImportUsage": false,
-+    "reportPrivateImportUsage": false
+    "reportPrivateImportUsage": false
 }

src/adaptive_graph_of_thoughts/validation.py

@@ -20,7 +20,7 @@ def validate_query_content(cls, v):
 
         for pattern in dangerous_patterns:
             if re.search(pattern, v, re.IGNORECASE | re.DOTALL):
-                raise ValueError(f"Query contains potentially dangerous content")
+                raise ValueError("Query contains potentially dangerous content")
 
         return v.strip()
 

tests/integration/api/test_mcp_inspector.py

@@ -9,34 +9,40 @@
 import requests
 from unittest.mock import patch
 
+
 class ProcessManager:
     """Manages background processes for HTTP server testing."""
-    
+
     def __init__(self):
         self.processes = []
-    
+
     def start_http_server(self, timeout=30):
         """Start HTTP server in background and wait for it to be ready."""
         print("Starting HTTP server for testing...")
-        
+
         # Change to project root directory
         project_root = Path(__file__).parent.parent.parent
-        
+
         # Start server process
         cmd = [
-            sys.executable, "-m", "uvicorn", 
+            sys.executable,
+            "-m",
+            "uvicorn",
             "src.adaptive_graph_of_thoughts.main:app",
-            "--host", "0.0.0.0",
-            "--port", "8000",
-            "--log-level", "info"
+            "--host",
+            "0.0.0.0",
+            "--port",
+            "8000",
+            "--log-level",
+            "info",
         ]
-        
+
         process = subprocess.Popen(
             cmd,
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             preexec_fn=os.setsid,  # Create new process group
-            cwd=project_root
+            cwd=project_root,
         )
         self.processes.append(process)
 
@@ -56,7 +62,7 @@ def start_http_server(self, timeout=30):
             time.sleep(2)
 
         raise TimeoutError("HTTP server failed to start within timeout")
-    
+
     def cleanup(self):
         """Clean up all managed processes."""
         for proc in self.processes:
@@ -69,22 +75,22 @@ def cleanup(self):
                 except ProcessLookupError:
                     pass
 
+
 def run_with_retry(command, max_retries=2, timeout=180):
     """Run command with retry logic and proper timeout handling."""
     for attempt in range(max_retries + 1):
         try:
             print(f"Attempt {attempt + 1}/{max_retries + 1}: {' '.join(command)}")
             result = subprocess.run(
-                command,
-                capture_output=True,
-                text=True,
-                timeout=timeout
+                command, capture_output=True, text=True, timeout=timeout
             )
             if result.returncode == 0:
                 print("✅ MCP Inspector validation passed!")
                 print("Output:", result.stdout)
                 return True
-            print(f"❌ Exit {result.returncode}. Stdout: {result.stdout} Stderr: {result.stderr}")
+            print(
+                f"❌ Exit {result.returncode}. Stdout: {result.stdout} Stderr: {result.stderr}"
+            )
             if attempt < max_retries:
                 wait = 5 * (attempt + 1)
                 print(f"⏳ Retrying in {wait}s...")
@@ -101,14 +107,13 @@ def run_with_retry(command, max_retries=2, timeout=180):
                 time.sleep(5)
     return False
 
+
 @pytest.fixture(autouse=True)
 def setup_test_environment():
     """Setup test environment and ensure dependencies are available."""
     try:
         res = subprocess.run(
-            ["mcp-inspector", "--version"],
-            capture_output=True,
-            timeout=10
+            ["mcp-inspector", "--version"], capture_output=True, timeout=10
         )
         if res.returncode != 0:
             pytest.skip("mcp-inspector not available")
@@ -120,21 +125,27 @@ def setup_test_environment():
     yield
     os.chdir(original_cwd)
 
+
 def test_inspector_http():
     """Test MCP Inspector with HTTP transport using improved process management."""
     pm = ProcessManager()
     try:
         pm.start_http_server(timeout=30)
         cmd = ["mcp-inspector", "-v", "validate", "http://localhost:8000/mcp"]
-        assert run_with_retry(cmd, max_retries=2, timeout=120), \
+        assert run_with_retry(cmd, max_retries=2, timeout=120), (
             "MCP Inspector HTTP validation failed after retries"
+        )
     except Exception as e:
         pytest.fail(f"HTTP server setup failed: {e}")
     finally:
         print("🧹 Cleaning up HTTP server...")
         pm.cleanup()
 
-@pytest.mark.parametrize("transport_mode", ["stdio"])
+
+@pytest.mark.parametrize(
+    "transport_mode",
+    ["stdio"],
+)
 def test_inspector_stdio(transport_mode):
     """Test MCP Inspector with STDIO transport using improved protocol handling."""
     os.chdir(Path(__file__).parent.parent.parent)
@@ -147,34 +158,44 @@ def test_inspector_stdio(transport_mode):
         sys.executable,
         "src/adaptive_graph_of_thoughts/main_stdio.py",
     ]
-    assert run_with_retry(cmd, max_retries=2, timeout=180), \
+    assert run_with_retry(cmd, max_retries=2, timeout=180), (
         f"MCP Inspector {transport_mode} validation failed after retries"
+    )
+
 
 def test_inspector_both_transports():
     """Test MCP Inspector with both HTTP and STDIO transports comprehensively."""
     pm = ProcessManager()
     results = {}
     try:
-        print("\n" + "="*50 + "\nTesting STDIO transport\n" + "="*50)
+        print("\n" + "=" * 50 + "\nTesting STDIO transport\n" + "=" * 50)
         stdio_cmd = [
-            "mcp-inspector", "-v", "validate", "stdio",
-            "--program", sys.executable, "src/adaptive_graph_of_thoughts/main_stdio.py",
+            "mcp-inspector",
+            "-v",
+            "validate",
+            "stdio",
+            "--program",
+            sys.executable,
+            "src/adaptive_graph_of_thoughts/main_stdio.py",
         ]
-        results['stdio'] = run_with_retry(stdio_cmd, max_retries=2, timeout=180)
+        results["stdio"] = run_with_retry(stdio_cmd, max_retries=2, timeout=180)
 
-        print("\n" + "="*50 + "\nTesting HTTP transport\n" + "="*50)
+        print("\n" + "=" * 50 + "\nTesting HTTP transport\n" + "=" * 50)
         pm.start_http_server(timeout=30)
         http_cmd = ["mcp-inspector", "-v", "validate", "http://localhost:8000/mcp"]
-        results['http'] = run_with_retry(http_cmd, max_retries=2, timeout=120)
+        results["http"] = run_with_retry(http_cmd, max_retries=2, timeout=120)
     except Exception as e:
         pytest.fail(f"Transport testing failed: {e}")
     finally:
         pm.cleanup()
 
-    print(f"\n📊 Test Results:\nSTDIO: {'✅ PASS' if results.get('stdio') else '❌ FAIL'}\n" +
-          f"HTTP: {'✅ PASS' if results.get('http') else '❌ FAIL'}")
-    assert results.get('stdio'), "STDIO transport validation failed"
-    assert results.get('http'), "HTTP transport validation failed"
+    print(
+        f"\n📊 Test Results:\nSTDIO: {'✅ PASS' if results.get('stdio') else '❌ FAIL'}\n"
+        + f"HTTP: {'✅ PASS' if results.get('http') else '❌ FAIL'}"
+    )
+    assert results.get("stdio"), "STDIO transport validation failed"
+    assert results.get("http"), "HTTP transport validation failed"
+
 
 def test_mcp_configuration_validity():
     """Test that MCP configuration is valid and includes necessary timeout settings."""
@@ -183,4 +204,4 @@ def test_mcp_configuration_validity():
     content = cfg.read_text()
     for key in ("mcp_stdio_timeout", "mcp_http_timeout", "mcp_inspector_retries"):
         assert key in content, f"Missing {key} setting"
-    print("✅ MCP configuration includes required timeout and retry settings")
+    print("✅ MCP configuration includes required timeout and retry settings")

tests/integration/api/test_stdio_mcp.py

@@ -77,7 +77,10 @@ def test_stdio_initialize(stdio_process):
     assert "mcp_version" in result_data and isinstance(result_data["mcp_version"], str)
 
 
-@pytest.mark.parametrize("query", ["test question"])
+@pytest.mark.parametrize(
+    "query",
+    ["test question"],
+)
 def test_stdio_call_tool(stdio_process, query):
     """Test calling the asr_got_query tool over STDIO."""
     request = {

tests/integration/stages/test_integration_initialization_stage.py

@@ -10,16 +10,19 @@
 from dataclasses import dataclass
 import sys
 import types
-from typing import Generator
+from typing import ClassVar, Generator
 
 class DummyDefaultParams:
-    initial_confidence = [0.9, 0.9, 0.9, 0.9]
-    initial_layer = "root_layer"
+    initial_confidence: ClassVar[list[float]] = [0.9, 0.9, 0.9, 0.9]
+    initial_layer: ClassVar[str] = "root_layer"
+
+
+DummyASRGot = type("obj", (), {"default_parameters": DummyDefaultParams()})
 
 
 @dataclass
 class DummySettings:
-    asr_got: type = type("obj", (), {"default_parameters": DummyDefaultParams()})
+    asr_got: type = DummyASRGot
 
 @pytest.fixture(scope="module", autouse=True)
 def stub_modules(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]:
@@ -582,7 +585,7 @@ async def test_initialization_stage_concurrent_execution_safety(settings_instanc
 
     # Execute multiple times concurrently
     tasks = []
-    for i in range(5):
+    for _ in range(5):
         session_data = GoTProcessorSessionData(query=query)
         tasks.append(stage.execute(current_session_data=session_data))
 

tests/test_security.py

@@ -1,6 +1,16 @@
 import pytest
 from fastapi.testclient import TestClient
 from src.adaptive_graph_of_thoughts.app_setup import create_app
+import re
+
+
+def validate_identifier(identifier: str, allowed_labels: set[str]) -> str:
+    """Validate Neo4j identifier to prevent injection attacks."""
+    if not re.match(r"^[A-Za-z][A-Za-z0-9_]*$", identifier):
+        raise ValueError("Invalid identifier")
+    if identifier not in allowed_labels:
+        raise ValueError("Identifier not allowed")
+    return identifier
 
 class TestSecurity:
     def test_authentication_required(self):

tests/unit/api/test_nlq.py

@@ -658,7 +658,7 @@ def test_nlq_endpoint_llm_query_logs_management(client, auth_headers, monkeypatc
 
     def tracking_llm(prompt: str) -> str:
         if "Convert" in prompt:
-            return f"MATCH (n) RETURN n // Query for test"
+            return "MATCH (n) RETURN n // Query for test"
         return "Test response"
 
     monkeypatch.setattr("adaptive_graph_of_thoughts.services.llm.ask_llm", tracking_llm)

tests/unit/config/test_config_validation.py

@@ -57,7 +57,10 @@ def test_minimal_config(minimal_config):
     assert "google_scholar" not in config_data_to_validate
 
 
-@pytest.mark.parametrize("missing_key", ["app", "asr_got", "mcp_settings"])
+@pytest.mark.parametrize(
+    "missing_key",
+    ["app", "asr_got", "mcp_settings"],
+)
 def test_missing_keys(tmp_path, base_config_dict, missing_key):
     cfg = copy.deepcopy(base_config_dict)
     cfg.pop(missing_key)
@@ -120,7 +123,10 @@ def test_unreadable_file(tmp_path):
         file_path.chmod(stat.S_IRUSR | stat.S_IWUSR)
 
 
-@pytest.mark.parametrize("port", [1, 65535])
+@pytest.mark.parametrize(
+    "port",
+    [1, 65535],
+)
 def test_boundary_values(tmp_path, port):
     cfg = copy.deepcopy(base_config_dict)
     cfg["app"]["port"] = port