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")


def test_cost_tracker_concurrent_record_usage() -> None:
    """Verify CostTracker.record_usage is thread-safe under concurrent calls.

    Spawns multiple threads that all call record_usage simultaneously.
    The final daily spend must equal the sum of all individual costs —
    no increments may be silently lost due to a race condition.
    """
    import threading

    from cleveragents.domain.models.core.cost_metadata import CostMetadata
    from cleveragents.providers.cost_tracker import CostTracker

    n_threads = 20
    input_tokens = 100
    output_tokens = 50
    provider = "openai"
    model = "gpt-4o"

    tracker = CostTracker()
    single_cost = tracker.estimate_cost(provider, model, input_tokens, output_tokens)
    expected_total = single_cost * n_threads

    errors: list[Exception] = []
    barrier = threading.Barrier(n_threads)

    def worker() -> None:
        meta = CostMetadata()
        try:
            barrier.wait()  # All threads start simultaneously
            tracker.record_usage(
                meta,
                provider=provider,
                model=model,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
            )
        except Exception as exc:
            errors.append(exc)

    threads = [threading.Thread(target=worker) for _ in range(n_threads)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    assert not errors, f"Threads raised exceptions: {errors}"

    actual = tracker.get_daily_spend()
    assert abs(actual - expected_total) < 1e-9, (
        f"Expected daily spend {expected_total:.6f}, got {actual:.6f} — "
        "possible race condition: cost increments were lost"
    )
    print("cost-tracker-concurrent-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,
        "cost-tracker-concurrent": test_cost_tracker_concurrent_record_usage,
    }
    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()