cleveragents-core/robot/helper_cost_controls.py

#!/usr/bin/env python3
"""Helper script for cost controls Robot Framework integration tests."""

from __future__ import annotations

import sys
from pathlib import Path

# Ensure local source is importable
_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)


def test_cost_entry_creation() -> None:
    """Verify CostEntry creation and estimation."""
    from cleveragents.providers.cost_table import CostEntry

    entry = CostEntry(input_cost_per_token=0.00001, output_cost_per_token=0.00003)
    cost = entry.estimate_cost(1000, 500)
    assert cost > 0.0, f"Expected positive cost, got {cost}"
    print("cost-entry-ok")


def test_cost_table_lookup() -> None:
    """Verify ProviderCostTable lookup."""
    from cleveragents.providers.cost_table import ProviderCostTable

    table = ProviderCostTable()
    entry = table.get_cost_entry("openai", "gpt-4o")
    assert entry.input_cost_per_token > 0
    providers = table.list_providers()
    assert "openai" in providers
    print("cost-table-ok")


def test_cost_tracker_budget() -> None:
    """Verify CostTracker budget enforcement."""
    from cleveragents.domain.models.core.cost_metadata import CostMetadata
    from cleveragents.providers.cost_tracker import BudgetStatus, CostTracker

    tracker = CostTracker(budget_per_plan=0.0001)
    metadata = CostMetadata()
    result = tracker.record_usage(
        metadata,
        provider="openai",
        model="gpt-4o",
        input_tokens=10000,
        output_tokens=5000,
    )
    # With a tiny budget, usage should exceed it
    assert result.status in (BudgetStatus.WARNING, BudgetStatus.EXCEEDED), (
        f"Expected warning/exceeded, got {result.status}"
    )
    print("cost-tracker-budget-ok")


def test_fallback_selector() -> None:
    """Verify FallbackSelector with no configured providers."""
    import os

    from cleveragents.config.settings import Settings
    from cleveragents.providers.fallback_selector import FallbackSelector
    from cleveragents.providers.registry import ProviderRegistry

    # Clear all API keys so no providers are configured
    api_key_vars = [
        "OPENAI_API_KEY",
        "ANTHROPIC_API_KEY",
        "GOOGLE_API_KEY",
        "AZURE_OPENAI_API_KEY",
        "AZURE_API_KEY",
        "OPENROUTER_API_KEY",
        "GEMINI_API_KEY",
        "GOOGLE_GEMINI_API_KEY",
        "GROQ_API_KEY",
        "TOGETHER_API_KEY",
        "COHERE_API_KEY",
        "PERPLEXITY_API_KEY",
    ]
    saved: dict[str, str] = {}
    for var in api_key_vars:
        val = os.environ.pop(var, None)
        if val is not None:
            saved[var] = val
    try:
        settings = Settings()
        registry = ProviderRegistry(settings)
        selector = FallbackSelector(registry=registry)
        result = selector.select()
        # With no API keys, no provider should be selected
        assert result.provider_type is None
        assert len(result.skipped) > 0
        print("fallback-selector-ok")
    finally:
        # Restore any keys we removed
        os.environ.update(saved)


def test_settings_defaults() -> None:
    """Verify Settings cost control defaults."""
    from cleveragents.config.settings import Settings

    settings = Settings()
    assert settings.budget_per_plan is None
    assert settings.budget_per_day is None
    assert settings.fallback_providers == []
    print("settings-defaults-ok")


def test_cost_metadata_display() -> None:
    """Verify CostMetadata display dict."""
    from cleveragents.domain.models.core.cost_metadata import CostMetadata

    meta = CostMetadata()
    meta.record_usage(
        input_tokens=100,
        output_tokens=50,
        cost=0.01,
        provider="openai",
    )
    display = meta.as_display_dict()
    assert "total_tokens" in display
    assert "total_cost_usd" in display
    print("cost-metadata-display-ok")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: helper_cost_controls.py <test-name>", file=sys.stderr)
        sys.exit(1)

    test_name = sys.argv[1]
    tests = {
        "cost-entry": test_cost_entry_creation,
        "cost-table": test_cost_table_lookup,
        "cost-tracker": test_cost_tracker_budget,
        "fallback-selector": test_fallback_selector,
        "settings-defaults": test_settings_defaults,
        "cost-metadata": test_cost_metadata_display,
    }
    test_fn = tests.get(test_name)
    if test_fn is None:
        print(f"Unknown test: {test_name}", file=sys.stderr)
        sys.exit(1)
    test_fn()