/tests/unit/test_langfuse.py

"""Tests for Langfuse observability integration.

Unit tests for the Langfuse event handler module (SDK v3 API).
"""

import os
import tempfile
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

from pentestgpt.core.events import Event, EventBus, EventType


@pytest.mark.unit
class TestLangfuseIntegration:
    """Tests for Langfuse integration."""

    def test_init_langfuse_disabled_via_flag(self):
        """Test that Langfuse is not initialized when disabled flag is True."""
        import importlib

        import pentestgpt.core.langfuse as langfuse_module

        importlib.reload(langfuse_module)

        result = langfuse_module.init_langfuse(disabled=True)
        assert result is False

    def test_init_langfuse_disabled_via_env_var(self):
        """Test that Langfuse is not initialized when LANGFUSE_ENABLED=false."""
        env_backup = os.environ.get("LANGFUSE_ENABLED")
        os.environ["LANGFUSE_ENABLED"] = "false"

        try:
            import importlib

            import pentestgpt.core.langfuse as langfuse_module

            importlib.reload(langfuse_module)

            result = langfuse_module.init_langfuse()
            assert result is False
        finally:
            if env_backup is not None:
                os.environ["LANGFUSE_ENABLED"] = env_backup
            else:
                os.environ.pop("LANGFUSE_ENABLED", None)

    def test_init_langfuse_with_mock_v3_api(self):
        """Test Langfuse initialization uses v3 get_client() API."""
        env_backup = os.environ.get("LANGFUSE_ENABLED")

        try:
            os.environ.pop("LANGFUSE_ENABLED", None)

            import importlib
            import sys

            import pentestgpt.core.langfuse as langfuse_module

            importlib.reload(langfuse_module)

            mock_langfuse_module = MagicMock()
            mock_client = MagicMock()
            mock_langfuse_module.get_client.return_value = mock_client

            with patch.dict(sys.modules, {"langfuse": mock_langfuse_module}):
                result = langfuse_module.init_langfuse()

            # Should use v3 get_client() API
            mock_langfuse_module.get_client.assert_called_once()
            assert result is True
            # Should have a user ID set
            assert langfuse_module._user_id is not None
        finally:
            if env_backup is not None:
                os.environ["LANGFUSE_ENABLED"] = env_backup
            else:
                os.environ.pop("LANGFUSE_ENABLED", None)

    def test_get_or_create_user_id_creates_new(self):
        """Test that a new user ID is created if none exists."""
        import importlib

        import pentestgpt.core.langfuse as langfuse_module

        importlib.reload(langfuse_module)

        with (
            tempfile.TemporaryDirectory() as tmpdir,
            patch.object(Path, "home", return_value=Path(tmpdir)),
        ):
            user_id = langfuse_module._get_or_create_user_id()

            # Should be a valid UUID format
            assert len(user_id) == 36
            assert user_id.count("-") == 4

            # File should exist
            user_id_file = Path(tmpdir) / ".pentestgpt" / "user_id"
            assert user_id_file.exists()
            assert user_id_file.read_text().strip() == user_id

    def test_get_or_create_user_id_reads_existing(self):
        """Test that existing user ID is read from file."""
        import importlib

        import pentestgpt.core.langfuse as langfuse_module

        importlib.reload(langfuse_module)

        with tempfile.TemporaryDirectory() as tmpdir:
            # Create existing user ID file
            pentestgpt_dir = Path(tmpdir) / ".pentestgpt"
            pentestgpt_dir.mkdir(parents=True)
            user_id_file = pentestgpt_dir / "user_id"
            existing_id = "existing-test-uuid-1234"
            user_id_file.write_text(existing_id)

            # Read the existing ID
            with patch.object(Path, "home", return_value=Path(tmpdir)):
                user_id = langfuse_module._get_or_create_user_id()
                assert user_id == existing_id

    def test_shutdown_langfuse_safe_when_not_initialized(self):
        """Test that shutdown_langfuse is safe to call when not initialized."""
        import importlib

        import pentestgpt.core.langfuse as langfuse_module

        importlib.reload(langfuse_module)

        # Should not raise any exception
        langfuse_module.shutdown_langfuse()

    def test_event_handlers_guard_against_no_client(self):
        """Test that event handlers safely do nothing when client is None."""
        import importlib

        import pentestgpt.core.langfuse as langfuse_module

        importlib.reload(langfuse_module)

        # Ensure client is None
        langfuse_module._langfuse_client = None
        langfuse_module._current_span = None

        # These should not raise exceptions
        langfuse_module._handle_state(Event(EventType.STATE_CHANGED, {"state": "running"}))
        langfuse_module._handle_message(Event(EventType.MESSAGE, {"text": "test"}))
        langfuse_module._handle_tool(Event(EventType.TOOL, {"status": "start", "name": "bash"}))
        langfuse_module._handle_flag(Event(EventType.FLAG_FOUND, {"flag": "test"}))

    def test_state_handler_creates_span_on_running(self):
        """Test that state handler creates a span when state is 'running' (v3 API)."""
        import importlib

        import pentestgpt.core.langfuse as langfuse_module

        importlib.reload(langfuse_module)

        mock_client = MagicMock()
        mock_span = MagicMock()
        mock_client.start_span.return_value = mock_span

        langfuse_module._langfuse_client = mock_client
        langfuse_module._current_span = None
        langfuse_module._user_id = "test-user-id-1234"

        langfuse_module._handle_state(
            Event(EventType.STATE_CHANGED, {"state": "running", "details": "10.10.11.234"})
        )

        # Verify start_span was called with correct parameters
        call_kwargs = mock_client.start_span.call_args.kwargs
        assert call_kwargs["name"] == "pentestgpt:10.10.11.234"
        assert call_kwargs["input"]["target"] == "10.10.11.234"
        assert call_kwargs["metadata"]["user_id"] == "test-user-id-1234"
        assert "session_id" in call_kwargs["metadata"]

        # Verify update_trace was called for user_id and session_id
        mock_span.update_trace.assert_called_once()
        trace_kwargs = mock_span.update_trace.call_args.kwargs
        assert trace_kwargs["user_id"] == "test-user-id-1234"
        assert "session_id" in trace_kwargs

        # Verify flush was called immediately after span creation
        assert mock_client.flush.called
        assert langfuse_module._current_span is mock_span

    def test_state_handler_ends_span_on_completed(self):
        """Test that state handler ends span and flushes on completion (v3 API)."""
        import importlib

        import pentestgpt.core.langfuse as langfuse_module

        importlib.reload(langfuse_module)

        mock_client = MagicMock()
        mock_span = MagicMock()

        langfuse_module._langfuse_client = mock_client
        langfuse_module._current_span = mock_span
        langfuse_module._session_target = "10.10.11.234"

        langfuse_module._handle_state(Event(EventType.STATE_CHANGED, {"state": "completed"}))

        mock_span.update.assert_called_once_with(
            output={"final_state": "completed", "target": "10.10.11.234"}
        )
        mock_span.end.assert_called_once()
        mock_client.flush.assert_called_once()
        assert langfuse_module._current_span is None
        assert langfuse_module._session_target is None

    def test_message_handler_creates_nested_span(self):
        """Test that message handler creates a nested span (v3 API)."""
        import importlib

        import pentestgpt.core.langfuse as langfuse_module

        importlib.reload(langfuse_module)

        mock_client = MagicMock()
        mock_span = MagicMock()
        mock_nested_span = MagicMock()
        mock_span.start_span.return_value = mock_nested_span

        langfuse_module._langfuse_client = mock_client
        langfuse_module._current_span = mock_span

        langfuse_module._handle_message(
            Event(EventType.MESSAGE, {"text": "Hello world", "type": "info"})
        )

        mock_span.start_span.assert_called_once_with(
            name="agent-message",
            input={"message_type": "info"},
            output={"text": "Hello world"},
        )
        mock_nested_span.end.assert_called_once()

    def test_tool_handler_creates_nested_span_on_start(self):
        """Test that tool handler creates a nested span on tool start (v3 API)."""
        import importlib

        import pentestgpt.core.langfuse as langfuse_module

        importlib.reload(langfuse_module)

        mock_client = MagicMock()
        mock_span = MagicMock()
        mock_nested_span = MagicMock()
        mock_span.start_span.return_value = mock_nested_span

        langfuse_module._langfuse_client = mock_client
        langfuse_module._current_span = mock_span

        langfuse_module._handle_tool(
            Event(EventType.TOOL, {"status": "start", "name": "bash", "args": {"cmd": "ls"}})
        )

        mock_span.start_span.assert_called_once_with(
            name="tool-bash",
            input={"cmd": "ls"},
            metadata={"tool_name": "bash"},
        )
        mock_nested_span.end.assert_called_once()

    def test_flag_handler_creates_nested_span(self):
        """Test that flag handler creates a nested span (v3 API)."""
        import importlib

        import pentestgpt.core.langfuse as langfuse_module

        importlib.reload(langfuse_module)

        mock_client = MagicMock()
        mock_span = MagicMock()
        mock_nested_span = MagicMock()
        mock_span.start_span.return_value = mock_nested_span

        langfuse_module._langfuse_client = mock_client
        langfuse_module._current_span = mock_span

        langfuse_module._handle_flag(
            Event(EventType.FLAG_FOUND, {"flag": "flag{test}", "context": "Found in output"})
        )

        mock_span.start_span.assert_called_once_with(
            name="flag-found",
            input={"context": "Found in output"},
            output={"flag": "flag{test}"},
            metadata={"flag": "flag{test}", "context": "Found in output"},
        )
        mock_nested_span.end.assert_called_once()

    def test_eventbus_integration(self):
        """Test that handlers are properly subscribed to EventBus (v3 API)."""
        import importlib

        import pentestgpt.core.langfuse as langfuse_module

        importlib.reload(langfuse_module)

        mock_client = MagicMock()
        mock_span = MagicMock()
        mock_client.start_span.return_value = mock_span

        langfuse_module._langfuse_client = mock_client
        langfuse_module._current_span = None
        langfuse_module._user_id = "test-user-id-1234"

        # Subscribe to events
        langfuse_module._subscribe_to_events()

        # Emit events via EventBus
        bus = EventBus.get()
        bus.emit_state("running", "Test session")

        # Verify span was created via EventBus subscription
        mock_client.start_span.assert_called()