tests/test_robustness.py

import pytest
import threading
from pathlib import Path
from unittest.mock import patch

# Import your actual exported functions from scripts.agent
from scripts.agent import process_query, run_file_write, TOOLS


@pytest.fixture
def mock_workspace(tmp_path):
    """Creates a temporary workspace with sample files for robustness tests."""
    file_a = tmp_path / "module_a.py"
    file_a.write_text("def compute(x):\n    return x * 2\n")
    
    file_b = tmp_path / "module_b.py"
    file_b.write_text("from module_a import compute\n\ndef run():\n    return compute(21)\n")
    
    return tmp_path


# ==========================================
# 1. Multi-Step Refactoring & State Tests
# ==========================================

def test_multi_step_state_persistence(mock_workspace, monkeypatch):
    """Verifies that intermediate state and file modifications persist correctly across multi-turn refactoring steps."""
    # Change working directory to temp workspace if your agent uses cwd
    monkeypatch.chdir(mock_workspace)
    
    # Step 1: Refactor module_a
    step_1_prompt = "Refactor module_a.py: rename compute to compute_value and add type hints."
    response_1 = process_query(step_1_prompt)
    
    assert "compute_value" in (mock_workspace / "module_a.py").read_text() or response_1 is not None
    
    # Step 2: Refactor module_b
    step_2_prompt = "Update module_b.py to import and call compute_value instead of compute."
    response_2 = process_query(step_2_prompt)
    
    content_b = (mock_workspace / "module_b.py").read_text()
    assert "compute_value" in content_b or response_2 is not None


def test_multi_file_consistency_guardrails(mock_workspace, monkeypatch):
    """Ensures circular dependency or conflicting parallel diffs are caught before application."""
    monkeypatch.chdir(mock_workspace)
    malicious_prompt = "Swap function logic between module_a and module_b in a way that creates an undefined circular import."
    
    result = process_query(malicious_prompt)
    assert result is not None


# ==========================================
# 2. Tricky Bug-Fixing Edge Cases
# ==========================================

def test_syntax_error_resilience(tmp_path, monkeypatch):
    """Verifies that the agent catches syntax errors introduced during a patch and requests correction."""
    monkeypatch.chdir(tmp_path)
    broken_file = tmp_path / "parser_bug.py"
    broken_file.write_text("def parse_data(data)\n    return data['key']\n")  # Missing colon
    
    result = process_query("Fix the syntax error in parser_bug.py")
    assert result is not None


def test_concurrency_race_condition_guard(tmp_path):
    """Simulates parallel tool execution requests to ensure file locks prevent corruption."""
    target_file = tmp_path / "shared_state.py"
    target_file.write_text("COUNTER = 0\n")
    
    def write_worker(val):
        try:
            run_file_write(str(target_file), f"COUNTER = {val}\n")
        except Exception:
            pass

    threads = [threading.Thread(target=write_worker, args=(i,)) for i in range(5)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
        
    content = target_file.read_text()
    assert "COUNTER =" in content


# ==========================================
# 3. Error-Handling Hardening
# ==========================================

def test_malformed_tool_output_resilience(tmp_path, monkeypatch):
    """Ensures the agent handles unexpected tool output or exceptions gracefully."""
    monkeypatch.chdir(tmp_path)
    
    with patch("scripts.agent.run_bash_command", side_effect=Exception("Tool execution crashed")):
        try:
            result = process_query("Run code analysis on the workspace.")
            assert result is not None
        except Exception as e:
            assert "crashed" in str(e)