test_workflow.py

import json
import time
from datetime import datetime, timedelta
from typing import Any
from unittest.mock import AsyncMock, Mock, patch

import pytest

# Mock Temporal and NATS modules before importing workflow code
with patch.dict(
    "sys.modules",
    {
        "temporalio": Mock(),
        "temporalio.workflow": Mock(),
        "temporalio.activity": Mock(),
        "temporalio.client": Mock(),
        "temporalio.worker": Mock(),
        "nats": Mock(),
        "nats.aio.client": Mock(),
        "nats.js": Mock(),
        "psycopg": Mock(),
        "orjson": Mock(),
        "slugify": Mock(),
        "prometheus_client": Mock(),
    },
):
    # Import workflow modules after mocking dependencies
    import os
    import sys

    # Add the apps directory to the path for imports
    sys.path.insert(0, os.path.join(os.path.dirname(__file__), "apps", "guard-codex"))
    sys.path.insert(0, os.path.join(os.path.dirname(__file__), "apps", "guard-api"))


class MockWorkflow:
    """Mock Temporal workflow decorator and utilities."""

    DEFAULT_VERSION = 1

    @staticmethod
    def defn(cls):
        return cls

    @staticmethod
    def run(func):
        return func

    @staticmethod
    def get_version(name, default, max_version):
        return default

    @staticmethod
    async def execute_activity(activity_func, *args, **kwargs):
        # Mock activity execution
        if hasattr(activity_func, "__name__"):
            if activity_func.__name__ == "extract_event_facts":
                return {"kind": "PR", "repo": "test/repo", "sha": "abc123"}
            elif activity_func.__name__ == "analyze_repo_state":
                return {"analysis": "complete"}
            elif activity_func.__name__ == "update_graph":
                return {"graph_updated": True}
            elif activity_func.__name__ == "render_docs":
                return "/tmp/docs"
            elif activity_func.__name__ == "publish_portal":
                return {"portal_url": "http://localhost:8080"}
            elif activity_func.__name__ == "cleanup_database":
                return {"cleaned_records": 100}
        return {}

    @staticmethod
    def now():
        return datetime.now()


class MockActivity:
    """Mock Temporal activity decorator."""

    @staticmethod
    def defn(func):
        return func


class MockTemporalClient:
    """Mock Temporal client."""

    @classmethod
    async def connect(cls, host, namespace=None):
        return cls()

    def create_workflow_handle(self, workflow_type, id, task_queue):
        handle = Mock()
        handle.start = AsyncMock()
        return handle


class MockNATSClient:
    """Mock NATS client."""

    def __init__(self):
        self.connected = False
        self.jetstream_instance = MockJetStream()

    async def connect(self, servers):
        self.connected = True

    def jetstream(self):
        return self.jetstream_instance

    async def close(self):
        self.connected = False


class MockJetStream:
    """Mock NATS JetStream."""

    async def subscribe(self, subject, cb, stream, durable, manual_ack=True):
        # Mock subscription
        pass


class MockMessage:
    """Mock NATS message."""

    def __init__(self, data):
        self.data = data

    async def ack(self):
        pass

    async def nak(self):
        pass


class MockWorker:
    """Mock Temporal worker."""

    def __init__(self, client, task_queue, workflows, activities, **kwargs):
        self.client = client
        self.task_queue = task_queue
        self.workflows = workflows
        self.activities = activities
        self.config = kwargs

    async def run(self):
        # Mock worker run
        pass


class MockMetrics:
    """Mock metrics functions."""

    @staticmethod
    def record_workflow_start(workflow_type):
        pass

    @staticmethod
    def record_workflow_completion(workflow_type, duration, success=True):
        pass

    @staticmethod
    def record_workflow_failure(workflow_type, failure_type):
        pass

    @staticmethod
    def record_codex_event(event_type, success=True):
        pass

    @staticmethod
    def start_metrics_server(port):
        pass


# Mock the workflow and activity modules
class MockCodexWorkflow:
    """Mock CodexWorkflow implementation."""

    async def run(self, event: dict[str, Any]):
        workflow_type = "CodexWorkflow"
        start_time = time.time()

        # Mock workflow execution
        event_type = event.get("event", "unknown")

        if event.get("should_fail"):
            raise Exception("Mock workflow failure")

        if event_type == "unknown":
            return {"skipped": True, "reason": "unknown_event"}

        return {"ok": True, "portal_url": "http://localhost:8080"}


class MockMaintenanceWorkflow:
    """Mock MaintenanceWorkflow implementation."""

    async def run(self) -> dict[str, Any]:
        return {
            "maintenance_completed": True,
            "cleanup_result": {"cleaned_records": 100},
            "timestamp": datetime.now(),
        }


# Mock activities
def mock_extract_event_facts(event):
    if event.get("event") == "pull_request":
        return {"kind": "PR", "repo": "test/repo", "sha": "abc123"}
    return {"kind": "unknown"}


def mock_analyze_repo_state(repo, sha):
    return {"analysis": "complete", "files_changed": 5}


def mock_update_graph(facts, analysis):
    return {"graph_updated": True, "nodes_added": 3}


def mock_render_docs(facts, analysis):
    return "/tmp/docs/output"


def mock_publish_portal(path):
    return {"portal_url": "http://localhost:8080", "docs_url": "http://localhost:8080/docs"}


def mock_cleanup_database():
    return {"cleaned_records": 100, "freed_space_mb": 50}


class TestCodexWorkflow:
    """Test suite for CodexWorkflow."""

    def setup_method(self):
        """Set up test fixtures."""
        self.workflow = MockCodexWorkflow()
        self.mock_metrics = MockMetrics()

    @pytest.mark.asyncio
    async def test_workflow_successful_execution(self):
        """Test successful workflow execution."""
        event = {
            "event": "pull_request",
            "action": "opened",
            "delivery_id": "12345",
            "pull_request": {"id": 1},
        }

        result = await self.workflow.run(event)

        assert result["ok"] is True
        assert "portal_url" in result

    @pytest.mark.asyncio
    async def test_workflow_unknown_event_handling(self):
        """Test workflow handling of unknown events."""
        event = {"event": "unknown", "delivery_id": "12345"}

        result = await self.workflow.run(event)

        assert result["skipped"] is True
        assert result["reason"] == "unknown_event"

    @pytest.mark.asyncio
    async def test_workflow_failure_handling(self):
        """Test workflow failure handling."""
        event = {"event": "pull_request", "should_fail": True, "delivery_id": "12345"}

        with pytest.raises(Exception, match="Mock workflow failure"):
            await self.workflow.run(event)

    def test_workflow_metrics_recording(self):
        """Test workflow metrics are recorded properly."""
        with patch("test_workflow.MockMetrics.record_workflow_start") as mock_start:
            with patch("test_workflow.MockMetrics.record_workflow_completion") as mock_completion:
                # Test that metrics functions can be called
                self.mock_metrics.record_workflow_start("CodexWorkflow")
                self.mock_metrics.record_workflow_completion("CodexWorkflow", 1.5, True)

                # Verify the mock functions were called
                assert mock_start.call_count >= 0  # Mock allows any number of calls
                assert mock_completion.call_count >= 0


class TestMaintenanceWorkflow:
    """Test suite for MaintenanceWorkflow."""

    def setup_method(self):
        """Set up test fixtures."""
        self.workflow = MockMaintenanceWorkflow()

    @pytest.mark.asyncio
    async def test_maintenance_workflow_execution(self):
        """Test maintenance workflow execution."""
        result = await self.workflow.run()

        assert result["maintenance_completed"] is True
        assert "cleanup_result" in result
        assert "timestamp" in result
        assert result["cleanup_result"]["cleaned_records"] == 100


class TestNATSIntegration:
    """Test suite for NATS messaging integration."""

    def setup_method(self):
        """Set up test fixtures."""
        self.nats_client = MockNATSClient()
        self.temporal_client = MockTemporalClient()

    @pytest.mark.asyncio
    async def test_nats_connection(self):
        """Test NATS connection establishment."""
        await self.nats_client.connect(["nats://localhost:4222"])
        assert self.nats_client.connected is True

    @pytest.mark.asyncio
    async def test_nats_jetstream_subscription(self):
        """Test NATS JetStream subscription."""
        js = self.nats_client.jetstream()

        # Mock subscription should not raise an exception
        await js.subscribe(
            "gh.*.*", cb=lambda msg: None, stream="GH", durable="CODEX", manual_ack=True
        )

    @pytest.mark.asyncio
    async def test_message_handler_success(self):
        """Test successful message handling."""
        event_data = {
            "event": "pull_request",
            "action": "opened",
            "delivery_id": "12345",
            "pull_request": {"id": 1},
        }

        message = MockMessage(json.dumps(event_data).encode())

        # Mock handler function
        async def handler(msg):
            try:
                evt = json.loads(msg.data.decode())
                wf_id = f"codex-{evt.get('delivery_id')}"

                handle = self.temporal_client.create_workflow_handle(
                    "CodexWorkflow", id=wf_id, task_queue="codex-task-queue"
                )
                await handle.start(evt)
                await msg.ack()
                return True
            except Exception:
                await msg.nak()
                return False

        result = await handler(message)
        assert result is True

    @pytest.mark.asyncio
    async def test_message_handler_invalid_json(self):
        """Test message handler with invalid JSON."""
        message = MockMessage(b"invalid json")

        async def handler(msg):
            try:
                json.loads(msg.data.decode())
                await msg.ack()
                return True
            except Exception:
                await msg.nak()
                return False

        result = await handler(message)
        assert result is False

    @pytest.mark.asyncio
    async def test_workflow_idempotency(self):
        """Test workflow idempotency with duplicate events."""
        event_data = {"event": "pull_request", "delivery_id": "12345", "pull_request": {"id": 1}}

        # Generate workflow ID the same way as the actual code
        wf_id = f"codex-{event_data.get('delivery_id')}"

        # First execution
        handle1 = self.temporal_client.create_workflow_handle(
            "CodexWorkflow", id=wf_id, task_queue="codex-task-queue"
        )
        await handle1.start(event_data)

        # Second execution with same ID should be idempotent
        handle2 = self.temporal_client.create_workflow_handle(
            "CodexWorkflow", id=wf_id, task_queue="codex-task-queue"
        )
        await handle2.start(event_data)

        # Both should succeed (mocked)
        assert handle1 is not None
        assert handle2 is not None


class TestTemporalWorker:
    """Test suite for Temporal worker configuration."""

    def setup_method(self):
        """Set up test fixtures."""
        self.client = MockTemporalClient()

    def test_worker_initialization(self):
        """Test worker initialization with proper configuration."""
        worker = MockWorker(
            self.client,
            task_queue="codex-task-queue",
            workflows=[MockCodexWorkflow, MockMaintenanceWorkflow],
            activities=[
                mock_extract_event_facts,
                mock_analyze_repo_state,
                mock_update_graph,
                mock_render_docs,
                mock_publish_portal,
                mock_cleanup_database,
            ],
            max_concurrent_activities=64,
            max_concurrent_workflow_tasks=256,
            max_concurrent_local_activities=32,
            activity_heartbeat_timeout=timedelta(seconds=30),
            activity_start_to_close_timeout=timedelta(minutes=10),
            activity_schedule_to_close_timeout=timedelta(minutes=15),
            workflow_task_timeout=timedelta(seconds=10),
            graceful_shutdown_timeout=timedelta(seconds=30),
        )

        assert worker.task_queue == "codex-task-queue"
        assert len(worker.workflows) == 2
        assert len(worker.activities) == 6
        assert worker.config["max_concurrent_activities"] == 64
        assert worker.config["max_concurrent_workflow_tasks"] == 256

    @pytest.mark.asyncio
    async def test_worker_run(self):
        """Test worker run method."""
        worker = MockWorker(
            self.client,
            task_queue="codex-task-queue",
            workflows=[MockCodexWorkflow],
            activities=[mock_extract_event_facts],
        )

        # Should not raise an exception
        await worker.run()


class TestWorkflowActivities:
    """Test suite for workflow activities."""

    def test_extract_event_facts(self):
        """Test event facts extraction."""
        pr_event = {"event": "pull_request", "action": "opened", "pull_request": {"id": 1}}

        result = mock_extract_event_facts(pr_event)
        assert result["kind"] == "PR"
        assert "repo" in result
        assert "sha" in result

        unknown_event = {"event": "unknown"}
        result = mock_extract_event_facts(unknown_event)
        assert result["kind"] == "unknown"

    def test_analyze_repo_state(self):
        """Test repository state analysis."""
        result = mock_analyze_repo_state("test/repo", "abc123")
        assert result["analysis"] == "complete"
        assert "files_changed" in result

    def test_update_graph(self):
        """Test knowledge graph update."""
        facts = {"kind": "PR", "repo": "test/repo"}
        analysis = {"files_changed": 5}

        result = mock_update_graph(facts, analysis)
        assert result["graph_updated"] is True
        assert "nodes_added" in result

    def test_render_docs(self):
        """Test documentation rendering."""
        facts = {"kind": "PR"}
        analysis = {"analysis": "complete"}

        result = mock_render_docs(facts, analysis)
        assert result.startswith("/tmp/docs")

    def test_publish_portal(self):
        """Test portal publishing."""
        result = mock_publish_portal("/tmp/docs/output")
        assert "portal_url" in result
        assert "docs_url" in result

    def test_cleanup_database(self):
        """Test database cleanup activity."""
        result = mock_cleanup_database()
        assert "cleaned_records" in result
        assert "freed_space_mb" in result
        assert result["cleaned_records"] > 0


class TestWorkflowIntegration:
    """Test suite for end-to-end workflow integration."""

    @pytest.mark.asyncio
    async def test_complete_workflow_pipeline(self):
        """Test complete workflow from NATS message to completion."""
        # Mock the complete pipeline
        nats_client = MockNATSClient()
        temporal_client = MockTemporalClient()
        workflow = MockCodexWorkflow()

        # Connect to NATS
        await nats_client.connect(["nats://localhost:4222"])

        # Create event
        event_data = {
            "event": "pull_request",
            "action": "opened",
            "delivery_id": "12345",
            "pull_request": {"id": 1},
        }

        # Process through workflow
        result = await workflow.run(event_data)

        assert result["ok"] is True
        assert "portal_url" in result

    def test_workflow_error_handling(self):
        """Test comprehensive error handling across the workflow."""
        # Test various error scenarios
        error_scenarios = [
            {"name": "invalid_json", "data": b"invalid json"},
            {"name": "missing_event", "data": json.dumps({}).encode()},
            {"name": "workflow_failure", "data": json.dumps({"should_fail": True}).encode()},
        ]

        for scenario in error_scenarios:
            message = MockMessage(scenario["data"])

            # Each scenario should be handled gracefully
            assert message is not None
            assert hasattr(message, "ack")
            assert hasattr(message, "nak")

    def test_metrics_integration(self):
        """Test metrics integration across workflow components."""
        metrics = MockMetrics()

        # Test all metrics functions are available
        metrics.record_workflow_start("CodexWorkflow")
        metrics.record_workflow_completion("CodexWorkflow", 1.5, True)
        metrics.record_workflow_failure("CodexWorkflow", "TestError")
        metrics.record_codex_event("pull_request", True)
        metrics.start_metrics_server(8090)

        # All should execute without error
        assert True


class TestWorkflowConfiguration:
    """Test suite for workflow configuration and environment setup."""

    def test_environment_variables(self):
        """Test environment variable handling."""
        # Test default values
        nats_url = os.getenv("NATS_URL", "nats://nats:4222")
        temporal_host = os.getenv("TEMPORAL_HOST", "temporal:7233")
        temporal_namespace = os.getenv("TEMPORAL_NAMESPACE", "default")
        metrics_port = int(os.getenv("METRICS_PORT", "8090"))

        assert nats_url is not None
        assert temporal_host is not None
        assert temporal_namespace is not None
        assert isinstance(metrics_port, int)

    def test_workflow_versioning(self):
        """Test workflow versioning support."""
        # Test version handling
        version = MockWorkflow.get_version("codex.analysis", MockWorkflow.DEFAULT_VERSION, 2)
        assert version == MockWorkflow.DEFAULT_VERSION

    def test_task_queue_configuration(self):
        """Test task queue configuration."""
        task_queue = "codex-task-queue"
        assert task_queue == "codex-task-queue"

    def test_timeout_configuration(self):
        """Test timeout configuration values."""
        timeouts = {
            "activity_heartbeat": timedelta(seconds=30),
            "activity_start_to_close": timedelta(minutes=10),
            "activity_schedule_to_close": timedelta(minutes=15),
            "workflow_task": timedelta(seconds=10),
            "graceful_shutdown": timedelta(seconds=30),
        }

        for name, timeout in timeouts.items():
            assert isinstance(timeout, timedelta)
            assert timeout.total_seconds() > 0


if __name__ == "__main__":
    pytest.main([__file__])