greydgl--pentestgpt/tests/unit/test_session.py

"""Tests for session management.

Unit tests for SessionInfo and SessionStore.
"""

from datetime import datetime
from pathlib import Path

import pytest

from pentestgpt.core.session import SessionInfo, SessionStatus, SessionStore


@pytest.mark.unit
class TestSessionInfo:
    """Tests for SessionInfo dataclass."""

    def test_create_session_info(self):
        """Test creating a SessionInfo instance."""
        session = SessionInfo(
            session_id="test123",
            target="10.10.11.234",
            created_at=datetime.now(),
        )
        assert session.session_id == "test123"
        assert session.target == "10.10.11.234"
        assert session.status == SessionStatus.RUNNING
        assert session.flags_found == []
        assert session.total_cost_usd == 0.0

    def test_session_to_dict(self):
        """Test serializing SessionInfo to dict."""
        now = datetime.now()
        session = SessionInfo(
            session_id="test123",
            target="10.10.11.234",
            created_at=now,
            status=SessionStatus.COMPLETED,
            task="Test task",
            model="claude-sonnet",
        )
        data = session.to_dict()
        assert data["session_id"] == "test123"
        assert data["target"] == "10.10.11.234"
        assert data["status"] == "completed"
        assert data["task"] == "Test task"
        assert data["model"] == "claude-sonnet"

    def test_session_from_dict(self):
        """Test deserializing SessionInfo from dict."""
        data = {
            "session_id": "test456",
            "target": "example.com",
            "created_at": "2024-01-01T12:00:00",
            "status": "paused",
            "task": "Solve CTF",
            "flags_found": [{"flag": "flag{test}", "context": "Found it"}],
            "total_cost_usd": 1.5,
            "model": "claude-opus",
        }
        session = SessionInfo.from_dict(data)
        assert session.session_id == "test456"
        assert session.target == "example.com"
        assert session.status == SessionStatus.PAUSED
        assert len(session.flags_found) == 1
        assert session.total_cost_usd == 1.5


@pytest.mark.unit
class TestSessionStore:
    """Tests for SessionStore."""

    @pytest.fixture
    def session_store(self, temp_sessions_dir: Path) -> SessionStore:
        """Create a SessionStore with temp directory."""
        return SessionStore(sessions_dir=temp_sessions_dir)

    def test_create_session(self, session_store: SessionStore):
        """Test creating a new session."""
        session = session_store.create(
            target="10.10.11.234",
            task="Solve CTF",
            model="claude-sonnet",
        )
        assert session.session_id is not None
        assert len(session.session_id) == 8
        assert session.target == "10.10.11.234"
        assert session.task == "Solve CTF"
        assert session_store.current == session

    def test_save_and_load_session(self, session_store: SessionStore, temp_sessions_dir: Path):
        """Test saving and loading a session."""
        session = session_store.create(
            target="example.com",
            task="Test task",
            model="claude-opus",
        )
        session_id = session.session_id

        # Verify file was created
        session_file = temp_sessions_dir / f"{session_id}.json"
        assert session_file.exists()

        # Load the session
        loaded = session_store.load(session_id)
        assert loaded is not None
        assert loaded.session_id == session_id
        assert loaded.target == "example.com"

    def test_list_sessions(self, session_store: SessionStore):
        """Test listing sessions."""
        # Create multiple sessions
        session_store.create("target1.com", "Task 1", "model1")
        session_store.create("target2.com", "Task 2", "model2")
        session_store.create("target1.com", "Task 3", "model3")

        # List all
        all_sessions = session_store.list_sessions()
        assert len(all_sessions) == 3

        # Filter by target
        target1_sessions = session_store.list_sessions("target1.com")
        assert len(target1_sessions) == 2

    def test_get_latest_session(self, session_store: SessionStore):
        """Test getting the most recent session."""
        session_store.create("example.com", "Task 1", "model1")
        session2 = session_store.create("example.com", "Task 2", "model2")

        latest = session_store.get_latest("example.com")
        assert latest is not None
        assert latest.session_id == session2.session_id

    def test_delete_session(self, session_store: SessionStore, temp_sessions_dir: Path):
        """Test deleting a session."""
        session = session_store.create("test.com", "Test", "model")
        session_id = session.session_id

        assert session_store.delete(session_id) is True
        assert not (temp_sessions_dir / f"{session_id}.json").exists()
        assert session_store.current is None

    def test_update_session_status(self, session_store: SessionStore):
        """Test updating session status."""
        session_store.create("test.com", "Test", "model")
        session_store.update_status(SessionStatus.COMPLETED)

        assert session_store.current is not None
        assert session_store.current.status == SessionStatus.COMPLETED

    def test_add_flag(self, session_store: SessionStore):
        """Test adding a flag to session."""
        session_store.create("test.com", "Test", "model")
        session_store.add_flag("flag{test123}", "Found in output")

        assert session_store.current is not None
        assert len(session_store.current.flags_found) == 1
        assert session_store.current.flags_found[0]["flag"] == "flag{test123}"

    def test_add_cost(self, session_store: SessionStore):
        """Test adding cost to session."""
        session_store.create("test.com", "Test", "model")
        session_store.add_cost(0.5)
        session_store.add_cost(0.3)

        assert session_store.current is not None
        assert session_store.current.total_cost_usd == pytest.approx(0.8)

    def test_load_nonexistent_session(self, session_store: SessionStore):
        """Test loading a session that doesn't exist."""
        result = session_store.load("nonexistent")
        assert result is None