cleveragents-core/robot/helper_estimation_actor.py

"""Helper utilities for estimation actor Robot integration tests.

Covers:
- EstimationResult domain model creation and serialization
- EstimationStubActor placeholder output
- Plan model estimation_result field
- Estimation invocation during Strategize→Execute transition
- ActorRole.ESTIMATOR enum value

Issue #890.
"""

from __future__ import annotations

import sys

from cleveragents.application.services.plan_executor import EstimationStubActor
from cleveragents.application.services.plan_lifecycle_service import (
    PlanLifecycleService,
)
from cleveragents.config.settings import Settings
from cleveragents.core.exceptions import ValidationError
from cleveragents.domain.models.acms.tiers import ActorRole
from cleveragents.domain.models.core.estimation import EstimationResult
from cleveragents.domain.models.core.plan import (
    PlanPhase,
    ProcessingState,
)


def _create_service() -> PlanLifecycleService:
    settings = Settings()
    return PlanLifecycleService(settings=settings)


def _estimation_result_model() -> None:
    """Test EstimationResult domain model creation and serialization."""
    # All fields
    result = EstimationResult(
        estimated_cost_usd=0.05,
        estimated_tokens=12000,
        estimated_steps=5,
        estimated_child_plans=2,
        estimated_time_seconds=30.0,
        risk_level="medium",
        risk_factors=("API rate limits", "Large codebase"),
        summary="Estimated 5 steps, ~$0.05 cost",
        raw_output="raw output",
    )
    assert result.estimated_cost_usd == 0.05
    assert result.estimated_tokens == 12000
    assert result.estimated_steps == 5
    assert result.estimated_child_plans == 2
    assert result.estimated_time_seconds == 30.0
    assert result.risk_level == "medium"
    assert len(result.risk_factors) == 2
    assert result.summary == "Estimated 5 steps, ~$0.05 cost"

    # Defaults
    default_result = EstimationResult()
    assert default_result.estimated_cost_usd is None
    assert default_result.estimated_tokens is None
    assert default_result.summary == ""
    assert default_result.risk_factors == ()

    # Serialization round-trip
    data = result.model_dump()
    restored = EstimationResult(**data)
    assert restored == result

    # Frozen check
    try:
        result.summary = "modified"  # type: ignore[misc]
        raise AssertionError("Should have raised ValidationError")
    except Exception:
        pass  # Expected: frozen model

    print("estimation-result-model-ok")


def _estimation_stub_actor() -> None:
    """Test EstimationStubActor placeholder output."""
    stub = EstimationStubActor()

    # Valid plan ID
    result = stub.estimate("01ARZ3NDEKTSV4RRFFQ69G5FAV")
    assert isinstance(result, EstimationResult)
    assert "stub" in result.summary.lower()

    # Empty plan ID should raise
    try:
        stub.estimate("")
        raise AssertionError("Should have raised ValidationError")
    except ValidationError:
        pass

    print("estimation-stub-actor-ok")


def _actor_role_estimator() -> None:
    """Test ActorRole.ESTIMATOR enum value."""
    assert hasattr(ActorRole, "ESTIMATOR")
    assert ActorRole.ESTIMATOR.value == "estimator"
    assert ActorRole("estimator") == ActorRole.ESTIMATOR
    print("actor-role-estimator-ok")


def _plan_estimation_field() -> None:
    """Test Plan model estimation_result field."""
    service = _create_service()

    # Create action with estimation actor
    service.create_action(
        name="local/est-robot-test",
        description="Robot test action",
        definition_of_done="Tests pass",
        strategy_actor="openai/gpt-4",
        execution_actor="openai/gpt-4",
        estimation_actor="local/estimator",
    )

    plan = service.use_action(action_name="local/est-robot-test")
    assert plan.estimation_actor == "local/estimator"
    assert plan.estimation_result is None

    # Set estimation result
    plan.estimation_result = EstimationResult(
        estimated_cost_usd=1.5,
        summary="Test estimation",
    )
    assert plan.estimation_result is not None
    assert plan.estimation_result.estimated_cost_usd == 1.5

    # CLI dict
    cli_dict = plan.as_cli_dict()
    assert "estimation" in cli_dict
    assert cli_dict["estimation"]["estimated_cost_usd"] == 1.5

    print("plan-estimation-field-ok")


def _estimation_lifecycle_integration() -> None:
    """Test estimation runs during Strategize→Execute transition."""
    service = _create_service()

    # Action WITH estimation actor
    service.create_action(
        name="local/est-lifecycle-robot",
        description="Lifecycle test",
        definition_of_done="Tests pass",
        strategy_actor="openai/gpt-4",
        execution_actor="openai/gpt-4",
        estimation_actor="local/estimator",
    )

    plan = service.use_action(action_name="local/est-lifecycle-robot")
    plan_id = plan.identity.plan_id

    # Drive through strategize
    service.start_strategize(plan_id)
    p = service.get_plan(plan_id)
    p.processing_state = ProcessingState.PROCESSING
    service.commit_plan(p)
    service.complete_strategize(plan_id)

    # Check the plan after complete_strategize — auto_progress might
    # have advanced it to Execute already (triggering estimation).
    p = service.get_plan(plan_id)
    if p.phase == PlanPhase.STRATEGIZE:
        service.execute_plan(plan_id)
        p = service.get_plan(plan_id)

    assert p.estimation_result is not None, (
        "estimation_result should be set after execute_plan"
    )
    assert "stub" in p.estimation_result.summary.lower()

    print("estimation-lifecycle-integration-ok")


def _no_estimation_lifecycle() -> None:
    """Test estimation is skipped when no estimation_actor is set."""
    service = _create_service()

    # Action WITHOUT estimation actor
    service.create_action(
        name="local/no-est-lifecycle-robot",
        description="No estimation test",
        definition_of_done="Tests pass",
        strategy_actor="openai/gpt-4",
        execution_actor="openai/gpt-4",
    )

    plan = service.use_action(action_name="local/no-est-lifecycle-robot")
    plan_id = plan.identity.plan_id

    # Drive through strategize
    service.start_strategize(plan_id)
    p = service.get_plan(plan_id)
    p.processing_state = ProcessingState.PROCESSING
    service.commit_plan(p)
    service.complete_strategize(plan_id)

    p = service.get_plan(plan_id)
    if p.phase == PlanPhase.STRATEGIZE:
        service.execute_plan(plan_id)
        p = service.get_plan(plan_id)

    assert p.estimation_result is None, (
        "estimation_result should be None when no estimation_actor is set"
    )

    print("no-estimation-lifecycle-ok")


_COMMANDS = {
    "estimation-result-model": _estimation_result_model,
    "estimation-stub-actor": _estimation_stub_actor,
    "actor-role-estimator": _actor_role_estimator,
    "plan-estimation-field": _plan_estimation_field,
    "estimation-lifecycle-integration": _estimation_lifecycle_integration,
    "no-estimation-lifecycle": _no_estimation_lifecycle,
}

if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1] not in _COMMANDS:
        print(f"Usage: {sys.argv[0]} <{'|'.join(_COMMANDS)}>")
        sys.exit(1)
    _COMMANDS[sys.argv[1]]()