项目文件夹

文件
2026-02-05 18:56:42 +03:00

173 行
6.2 KiB
Python

"""Tests for session management.
Unit tests for SessionInfo and SessionStore.
"""
from datetime import datetime
from pathlib import Path
import pytest
from excalibur.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