cleveragents-core/features/steps/cost_tracker_steps.py

"""Step definitions for CostTracker service tests."""

from __future__ import annotations

import tempfile
from pathlib import Path

from behave import given, then, when

from cleveragents.infrastructure.database.cost_tracker import CostRecord, CostTracker


@given("a temporary data directory for cost tracking")
def step_create_temp_dir(context):
    """Create a temporary directory for cost tracking."""
    context.temp_dir = tempfile.mkdtemp()
    context.data_dir = Path(context.temp_dir)


@given("a CostTracker instance with default model pricing")
def step_create_cost_tracker(context):
    """Create a CostTracker instance with default pricing."""
    context.cost_tracker = CostTracker(context.data_dir)
    context.recorded_records = []


@when('I record usage for session "{session_id}" and plan "{plan_id}"')
def step_record_usage(context, session_id, plan_id):
    """Record usage for a session and plan."""
    for row in context.table:
        tokens_in = int(row["tokens_in"])
        tokens_out = int(row["tokens_out"])
        model = row["model"]

        record = context.cost_tracker.record_usage(
            session_id=session_id,
            plan_id=plan_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            model=model,
        )
        context.recorded_records.append(record)


@then("the recorded cost should be approximately {cost} USD")
def step_check_recorded_cost(context, cost):
    """Check the recorded cost."""
    expected_cost = float(cost)
    actual_cost = context.recorded_records[-1].cost_usd
    assert abs(actual_cost - expected_cost) < 0.000001, (
        f"Expected {expected_cost}, got {actual_cost}"
    )


@then('the cost record should have session_id "{session_id}"')
def step_check_session_id(context, session_id):
    """Check the session_id of the recorded cost."""
    assert context.recorded_records[-1].session_id == session_id


@then('the cost record should have plan_id "{plan_id}"')
def step_check_plan_id(context, plan_id):
    """Check the plan_id of the recorded cost."""
    assert context.recorded_records[-1].plan_id == plan_id


@when('I record multiple usages for session "{session_id}"')
def step_record_multiple_usages_session(context, session_id):
    """Record multiple usages for a session."""
    for row in context.table:
        plan_id = row["plan_id"]
        tokens_in = int(row["tokens_in"])
        tokens_out = int(row["tokens_out"])
        model = row["model"]

        record = context.cost_tracker.record_usage(
            session_id=session_id,
            plan_id=plan_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            model=model,
        )
        context.recorded_records.append(record)


@then('the total session cost for "{session_id}" should be approximately {cost} USD')
def step_check_session_cost(context, session_id, cost):
    """Check the total cost for a session."""
    expected_cost = float(cost)
    actual_cost = context.cost_tracker.get_session_cost(session_id)
    assert abs(actual_cost - expected_cost) < 0.000001, (
        f"Expected {expected_cost}, got {actual_cost}"
    )


@when('I record multiple usages for plan "{plan_id}"')
def step_record_multiple_usages_plan(context, plan_id):
    """Record multiple usages for a plan."""
    for row in context.table:
        session_id = row["session_id"]
        tokens_in = int(row["tokens_in"])
        tokens_out = int(row["tokens_out"])
        model = row["model"]

        record = context.cost_tracker.record_usage(
            session_id=session_id,
            plan_id=plan_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            model=model,
        )
        context.recorded_records.append(record)


@then('the total plan cost for "{plan_id}" should be approximately {cost} USD')
def step_check_plan_cost(context, plan_id, cost):
    """Check the total cost for a plan."""
    expected_cost = float(cost)
    actual_cost = context.cost_tracker.get_plan_cost(plan_id)
    assert abs(actual_cost - expected_cost) < 0.000001, (
        f"Expected {expected_cost}, got {actual_cost}"
    )


@when("I create a new CostTracker instance with the same data directory")
def step_create_new_tracker(context):
    """Create a new CostTracker instance with the same data directory."""
    context.cost_tracker = CostTracker(context.data_dir)


@given("a CostTracker instance with custom model pricing")
def step_create_tracker_custom_pricing(context):
    """Create a CostTracker instance with custom pricing."""
    pricing = {}
    for row in context.table:
        model = row["model"]
        input_price = float(row["input_price"])
        output_price = float(row["output_price"])
        pricing[model] = {"input": input_price, "output": output_price}

    context.cost_tracker = CostTracker(context.data_dir, model_pricing=pricing)
    context.recorded_records = []


@then('I should get {count} entries for session "{session_id}"')
def step_check_session_entries_count(context, count, session_id):
    """Check the number of entries for a session."""
    entries = context.cost_tracker.get_session_entries(session_id)
    assert len(entries) == int(count), f"Expected {count} entries, got {len(entries)}"
    context.session_entries = entries


@then("each entry should have the correct session_id and plan_id")
def step_check_entries_ids(context):
    """Check that each entry has the correct IDs."""
    entries = getattr(context, "session_entries", None) or getattr(
        context, "plan_entries", []
    )
    for entry in entries:
        assert isinstance(entry, CostRecord)
        assert entry.session_id is not None
        assert entry.plan_id is not None


@then('I should get {count} entries for plan "{plan_id}"')
def step_check_plan_entries_count(context, count, plan_id):
    """Check the number of entries for a plan."""
    entries = context.cost_tracker.get_plan_entries(plan_id)
    assert len(entries) == int(count), f"Expected {count} entries, got {len(entries)}"
    context.plan_entries = entries