cleveragents-core/features/steps/session_cli_steps.py

"""Step definitions for the Session CLI feature."""

from __future__ import annotations

import json
import os
import tempfile
from datetime import datetime
from typing import Any
from unittest.mock import MagicMock

from behave import given, then, when
from behave.runner import Context
from typer.testing import CliRunner
from ulid import ULID

from cleveragents.cli.commands import session as session_mod
from cleveragents.cli.commands.session import app as session_app
from cleveragents.domain.models.core.session import (
    MessageRole,
    Session,
    SessionImportError,
    SessionMessage,
    SessionNotFoundError,
    SessionTokenUsage,
)

_SESSION_ID = str(ULID())
_SESSION_ID_2 = str(ULID())

_ENVELOPE_KEYS = {"command", "status", "exit_code", "data", "timing", "messages"}


def _unwrap_envelope(parsed: Any) -> Any:
    """Return the ``data`` field if *parsed* is a spec envelope, else *parsed* as-is."""
    if isinstance(parsed, dict) and _ENVELOPE_KEYS.issubset(parsed.keys()):
        return parsed["data"]
    return parsed


def _make_session(
    *,
    session_id: str | None = None,
    actor_name: str | None = None,
    messages: list[SessionMessage] | None = None,
) -> Session:
    """Create a test Session instance."""
    return Session(
        session_id=session_id or str(ULID()),
        actor_name=actor_name,
        namespace="local",
        messages=messages or [],
        token_usage=SessionTokenUsage(
            input_tokens=100,
            output_tokens=50,
            estimated_cost=0.005,
        ),
        created_at=datetime.now(),
        updated_at=datetime.now(),
    )


def _make_message(
    role: MessageRole = MessageRole.USER,
    content: str = "Hello",
    sequence: int = 0,
) -> SessionMessage:
    """Create a test SessionMessage."""
    return SessionMessage(
        message_id=str(ULID()),
        role=role,
        content=content,
        sequence=sequence,
        timestamp=datetime.now(),
    )


def _cleanup(context: Context) -> None:
    """Clean up temporary files."""
    for path in getattr(context, "_cleanup_paths", []):
        if os.path.exists(path):
            os.unlink(path)


# ---------------------------------------------------------------------------
# Background
# ---------------------------------------------------------------------------


@given("a session CLI runner with mocked service")
def step_session_cli_runner(context: Context) -> None:
    """Set up CLI runner with a mocked session service."""
    context.runner = CliRunner()
    context.mock_service = MagicMock()
    context._cleanup_paths: list[str] = []

    # Default: create returns a session
    default_session = _make_session(session_id=_SESSION_ID)
    context.mock_service.create.return_value = default_session

    # Patch the module-level service
    session_mod._service = context.mock_service

    def cleanup() -> None:
        session_mod._service = None
        _cleanup(context)

    context.add_cleanup(cleanup)


# ---------------------------------------------------------------------------
# Create
# ---------------------------------------------------------------------------


@when("I run session CLI create with no arguments")
def step_create_no_args(context: Context) -> None:
    context.result = context.runner.invoke(session_app, ["create"])


@when('I run session CLI create with --actor "{actor}"')
def step_create_with_actor(context: Context, actor: str) -> None:
    session = _make_session(session_id=_SESSION_ID, actor_name=actor)
    context.mock_service.create.return_value = session
    context.result = context.runner.invoke(session_app, ["create", "--actor", actor])


@when("I run session CLI create with --format json")
def step_create_json(context: Context) -> None:
    context.result = context.runner.invoke(session_app, ["create", "--format", "json"])


@then("the session CLI create should succeed")
def step_create_succeeds(context: Context) -> None:
    assert context.result.exit_code == 0, (
        f"Expected exit code 0, got {context.result.exit_code}: {context.result.output}"
    )


# ---------------------------------------------------------------------------
# List
# ---------------------------------------------------------------------------


@given("there are no mocked sessions")
def step_no_sessions(context: Context) -> None:
    context.mock_service.list.return_value = []


@given("there are mocked existing sessions")
def step_existing_sessions(context: Context) -> None:
    sessions = [
        _make_session(
            session_id=_SESSION_ID,
            actor_name="openai/gpt-4",
            messages=[_make_message(sequence=0)],
        ),
        _make_session(session_id=_SESSION_ID_2),
    ]
    context.mock_service.list.return_value = sessions


@when("I run session CLI list")
def step_list(context: Context) -> None:
    context.result = context.runner.invoke(session_app, ["list"])


@when("I run session CLI list with --format json")
def step_list_json(context: Context) -> None:
    context.result = context.runner.invoke(session_app, ["list", "--format", "json"])


@then("the session CLI should show all sessions in a table")
def step_list_shows_table(context: Context) -> None:
    assert context.result.exit_code == 0
    assert "Sessions" in context.result.output


# ---------------------------------------------------------------------------
# Show
# ---------------------------------------------------------------------------


@given("there is a mocked session with messages")
def step_session_with_messages(context: Context) -> None:
    messages = [
        _make_message(MessageRole.USER, "Hello", 0),
        _make_message(MessageRole.ASSISTANT, "Hi there", 1),
    ]
    session = _make_session(
        session_id=_SESSION_ID,
        actor_name="openai/gpt-4",
        messages=messages,
    )
    context.mock_service.get.return_value = session
    context.session_id = _SESSION_ID


@when("I run session CLI show with a valid session ID")
def step_show_valid(context: Context) -> None:
    context.result = context.runner.invoke(session_app, ["show", context.session_id])


@when("I run session CLI show with --format json")
def step_show_json(context: Context) -> None:
    context.result = context.runner.invoke(
        session_app, ["show", context.session_id, "--format", "json"]
    )


@when("I run session CLI show with an invalid session ID")
def step_show_invalid(context: Context) -> None:
    context.mock_service.get.side_effect = SessionNotFoundError(
        "Session 'INVALID' not found"
    )
    context.result = context.runner.invoke(session_app, ["show", "INVALID"])


@then("the session CLI show should succeed")
def step_show_succeeds(context: Context) -> None:
    assert context.result.exit_code == 0, (
        f"Expected exit code 0, got {context.result.exit_code}: {context.result.output}"
    )


# ---------------------------------------------------------------------------
# Delete
# ---------------------------------------------------------------------------


@given("there is a mocked session to delete")
def step_session_to_delete(context: Context) -> None:
    session = _make_session(session_id=_SESSION_ID)
    context.mock_service.get.return_value = session
    context.mock_service.delete.return_value = None
    context.session_id = _SESSION_ID


@when("I run session CLI delete with --yes")
def step_delete_yes(context: Context) -> None:
    context.result = context.runner.invoke(
        session_app, ["delete", context.session_id, "--yes"]
    )


@when("I run session CLI delete with a non-existent ID")
def step_delete_nonexistent(context: Context) -> None:
    context.mock_service.get.side_effect = SessionNotFoundError("Session not found")
    context.result = context.runner.invoke(
        session_app, ["delete", "NONEXISTENT", "--yes"]
    )


@then("the session CLI delete should succeed")
def step_delete_succeeds(context: Context) -> None:
    assert context.result.exit_code == 0, (
        f"Expected exit code 0, got {context.result.exit_code}: {context.result.output}"
    )


# ---------------------------------------------------------------------------
# Export
# ---------------------------------------------------------------------------


@given("there is a mocked session for export")
def step_session_for_export(context: Context) -> None:
    session = _make_session(session_id=_SESSION_ID)
    context.mock_service.get.return_value = session
    context.mock_service.export_session.return_value = session.as_export_dict()
    context.session_id = _SESSION_ID


@given("there is an existing export file")
def step_existing_export_file(context: Context) -> None:
    fd, path = tempfile.mkstemp(suffix=".json")
    with os.fdopen(fd, "w") as fh:
        fh.write("{}")
    context.existing_export_path = path
    context._cleanup_paths.append(path)


@when("I run session CLI export with no output file")
def step_export_stdout(context: Context) -> None:
    context.result = context.runner.invoke(session_app, ["export", context.session_id])


@when("I run session CLI export with --output to a temp file")
def step_export_to_file(context: Context) -> None:
    fd, path = tempfile.mkstemp(suffix=".json")
    os.close(fd)
    os.unlink(path)  # Remove so export can create it
    context.export_path = path
    context._cleanup_paths.append(path)
    context.result = context.runner.invoke(
        session_app, ["export", context.session_id, "--output", path]
    )


@when("I run session CLI export to an existing file without --force")
def step_export_no_force(context: Context) -> None:
    context.result = context.runner.invoke(
        session_app,
        ["export", context.session_id, "--output", context.existing_export_path],
    )


@when("I run session CLI export to an existing file with --force")
def step_export_with_force(context: Context) -> None:
    context.result = context.runner.invoke(
        session_app,
        [
            "export",
            context.session_id,
            "--output",
            context.existing_export_path,
            "--force",
        ],
    )


@when("I run session CLI export with a non-existent session ID")
def step_export_nonexistent(context: Context) -> None:
    context.mock_service.export_session.side_effect = SessionNotFoundError(
        "Session not found"
    )
    context.result = context.runner.invoke(session_app, ["export", "NONEXISTENT"])


@then("the session CLI export should succeed")
def step_export_succeeds(context: Context) -> None:
    assert context.result.exit_code == 0, (
        f"Expected exit code 0, got {context.result.exit_code}: {context.result.output}"
    )


@then("the exported file should exist")
def step_exported_file_exists(context: Context) -> None:
    assert os.path.exists(context.export_path), (
        f"Export file not found: {context.export_path}"
    )


# ---------------------------------------------------------------------------
# Import
# ---------------------------------------------------------------------------


@given("there is a valid session export file")
def step_valid_export_file(context: Context) -> None:
    session = _make_session(session_id=_SESSION_ID, actor_name="openai/gpt-4")
    export_data = session.as_export_dict()
    fd, path = tempfile.mkstemp(suffix=".json")
    with os.fdopen(fd, "w") as fh:
        json.dump(export_data, fh, default=str)
    context.import_path = path
    context._cleanup_paths.append(path)

    imported = _make_session(session_id=str(ULID()), actor_name="openai/gpt-4")
    context.mock_service.import_session.return_value = imported


@given("there is a corrupt session export file")
def step_corrupt_export_file(context: Context) -> None:
    fd, path = tempfile.mkstemp(suffix=".json")
    with os.fdopen(fd, "w") as fh:
        json.dump({"schema_version": "1.0", "bad": "data"}, fh)
    context.corrupt_path = path
    context._cleanup_paths.append(path)
    context.mock_service.import_session.side_effect = SessionImportError(
        "Checksum verification failed"
    )


@given("there is an invalid JSON file")
def step_invalid_json_file(context: Context) -> None:
    fd, path = tempfile.mkstemp(suffix=".json")
    with os.fdopen(fd, "w") as fh:
        fh.write("{not valid json")
    context.invalid_json_path = path
    context._cleanup_paths.append(path)


@when("I run session CLI import with the export file")
def step_import_valid(context: Context) -> None:
    context.result = context.runner.invoke(
        session_app, ["import", "--input", context.import_path]
    )


@when("I run session CLI import with a non-existent file")
def step_import_nonexistent(context: Context) -> None:
    context.result = context.runner.invoke(
        session_app, ["import", "--input", "/tmp/nonexistent_session_file.json"]
    )


@when("I run session CLI import with the corrupt file")
def step_import_corrupt(context: Context) -> None:
    context.result = context.runner.invoke(
        session_app, ["import", "--input", context.corrupt_path]
    )


@when("I run session CLI import with the invalid JSON file")
def step_import_invalid_json(context: Context) -> None:
    context.result = context.runner.invoke(
        session_app, ["import", "--input", context.invalid_json_path]
    )


@then("the session CLI import should succeed")
def step_import_succeeds(context: Context) -> None:
    assert context.result.exit_code == 0, (
        f"Expected exit code 0, got {context.result.exit_code}: {context.result.output}"
    )


# ---------------------------------------------------------------------------
# Tell
# ---------------------------------------------------------------------------


@given("there is a mocked session for tell")
def step_session_for_tell(context: Context) -> None:
    session = _make_session(session_id=_SESSION_ID)
    context.mock_service.get.return_value = session
    context.mock_service.append_message.side_effect = [
        _make_message(MessageRole.USER, "Hello, world", 0),
        _make_message(MessageRole.ASSISTANT, "Acknowledged: Hello, world", 1),
    ]
    context.session_id = _SESSION_ID


@when('I run session CLI tell with a prompt "{prompt}"')
def step_tell_prompt(context: Context, prompt: str) -> None:
    context.mock_service.append_message.side_effect = [
        _make_message(MessageRole.USER, prompt, 0),
        _make_message(MessageRole.ASSISTANT, f"Acknowledged: {prompt}", 1),
    ]
    context.result = context.runner.invoke(
        session_app,
        ["tell", "--session", context.session_id, prompt],
    )


@when('I run session CLI tell with --actor "{actor}" and prompt "{prompt}"')
def step_tell_with_actor(context: Context, actor: str, prompt: str) -> None:
    context.mock_service.append_message.side_effect = [
        _make_message(MessageRole.USER, prompt, 0),
        _make_message(
            MessageRole.ASSISTANT,
            f"[{actor}] Acknowledged: {prompt}",
            1,
        ),
    ]
    context.result = context.runner.invoke(
        session_app,
        ["tell", "--session", context.session_id, "--actor", actor, prompt],
    )


@when("I run session CLI tell to a non-existent session")
def step_tell_nonexistent(context: Context) -> None:
    context.mock_service.append_message.side_effect = SessionNotFoundError(
        "Session not found"
    )
    context.result = context.runner.invoke(
        session_app,
        ["tell", "--session", "NONEXISTENT", "Hello"],
    )


@then("the session CLI tell should succeed")
def step_tell_succeeds(context: Context) -> None:
    assert context.result.exit_code == 0, (
        f"Expected exit code 0, got {context.result.exit_code}: {context.result.output}"
    )


# ---------------------------------------------------------------------------
# Generic assertions
# ---------------------------------------------------------------------------


@then('the session CLI output should contain "{text}"')
def step_output_contains(context: Context, text: str) -> None:
    assert text in context.result.output, (
        f"Expected '{text}' in output:\n{context.result.output}"
    )


@then("the session CLI output should be valid JSON")
def step_output_valid_json(context: Context) -> None:
    try:
        json.loads(context.result.output)
    except json.JSONDecodeError as exc:
        raise AssertionError(
            f"Output is not valid JSON:\n{context.result.output}"
        ) from exc


@then('the session CLI JSON should contain "{key}"')
def step_json_contains_key(context: Context, key: str) -> None:
    parsed = json.loads(context.result.output)
    data = _unwrap_envelope(parsed)
    assert key in data, f"Key '{key}' not found in JSON: {data}"


@then("the session CLI should exit with error")
def step_exit_with_error(context: Context) -> None:
    assert context.result.exit_code != 0, (
        f"Expected non-zero exit code, got {context.result.exit_code}: "
        f"{context.result.output}"
    )