lambdagent/tests/test_agentpaas.py

"""Tests for agentpaas core components."""
import pytest
from agentpaas.config import AgentPaaSConfig
from agentpaas.engine.cache import CompileCache
from agentpaas.engine.retry import RetryPolicy, CircuitBreaker
from agentpaas.billing.pricing import calculate_cost, PRICING
from agentpaas.observability.metrics import Metrics
from agentpaas.tenant.rbac import check_permission
from agentpaas.api.middleware.auth import TenantContext, hash_key, verify_key_hash


# Config

def test_config_defaults():
    cfg = AgentPaaSConfig()
    assert cfg.port == 8000
    assert "sqlite" in cfg.database_url


# CompileCache

def test_cache_hit_miss():
    cache = CompileCache(maxsize=3)
    cache.put("h1", "term1")
    assert cache.get("h1") == "term1"
    assert cache.get("h2") is None
    assert cache.hits == 1
    assert cache.misses == 1


def test_cache_lru():
    cache = CompileCache(maxsize=2)
    cache.put("a", 1)
    cache.put("b", 2)
    cache.put("c", 3)  # evicts "a"
    assert cache.get("a") is None
    assert cache.get("c") == 3


# RetryPolicy

def test_retry_should_retry():
    rp = RetryPolicy(max_attempts=3)
    assert rp.should_retry(TimeoutError("timeout"), 0) == True
    assert rp.should_retry(TimeoutError("timeout"), 3) == False
    assert rp.should_retry(ValueError("bad"), 0) == False


def test_retry_delay():
    rp = RetryPolicy(base_delay=1.0, exponential_base=2.0, max_delay=10.0)
    assert rp.get_delay(0) == 1.0
    assert rp.get_delay(1) == 2.0
    assert rp.get_delay(2) == 4.0
    assert rp.get_delay(10) == 10.0  # capped


# CircuitBreaker

def test_circuit_breaker_flow():
    cb = CircuitBreaker(failure_threshold=3, recovery_timeout=0.1)
    assert cb.allow_request() == True
    cb.record_failure()
    cb.record_failure()
    assert cb.allow_request() == True  # not yet at threshold
    cb.record_failure()
    assert cb.state == CircuitBreaker.OPEN
    assert cb.allow_request() == False


def test_circuit_breaker_recovery():
    import time
    cb = CircuitBreaker(failure_threshold=2, recovery_timeout=0.1)
    cb.record_failure()
    cb.record_failure()
    assert cb.state == CircuitBreaker.OPEN
    time.sleep(0.15)
    assert cb.allow_request() == True  # half-open
    cb.record_success()
    assert cb.state == CircuitBreaker.CLOSED


# Billing

def test_pricing():
    cost = calculate_cost("gpt-4o", 1000, 500)
    assert cost > 0
    assert cost == (1000 / 1_000_000 * 2.5 + 500 / 1_000_000 * 10.0)


def test_pricing_unknown_model():
    cost = calculate_cost("unknown-model", 1000, 1000)
    assert cost > 0  # uses default pricing


# Metrics

def test_metrics_counter():
    m = Metrics()
    m.inc("requests", {"agent": "a1"})
    m.inc("requests", {"agent": "a1"})
    snap = m.snapshot()
    assert snap["counters"]["requests{agent=a1}"] == 2


def test_metrics_observe():
    m = Metrics()
    m.observe("latency", 100.0)
    m.observe("latency", 200.0)
    snap = m.snapshot()
    assert snap["histograms"]["latency"]["mean"] == 150.0


# RBAC

def test_rbac_admin():
    ctx = TenantContext("t1", "org", "pro", role="admin")
    assert check_permission(ctx, "tenants:read") == True
    assert check_permission(ctx, "agents:execute") == True


def test_rbac_viewer():
    ctx = TenantContext("t1", "org", "pro", role="viewer")
    assert check_permission(ctx, "agents:read") == True
    assert check_permission(ctx, "agents:write") == False


def test_rbac_scopes():
    ctx = TenantContext("t1", "org", "pro", role="viewer", scopes=["agents:execute"])
    assert check_permission(ctx, "agents:execute") == True


# Auth

def test_hash_key():
    # FIX-08: hash_key now uses PBKDF2 with a random salt per call.
    # Format: "pbkdf2:<32-hex salt>:<64-hex hash>" — 7 + 32 + 1 + 64 = 104 chars.
    h1 = hash_key("ap_test123")
    h2 = hash_key("ap_test123")
    assert h1.startswith("pbkdf2:")
    assert len(h1) == 104
    # Random salt → two calls on the same input must NOT collide.
    assert h1 != h2, "random-salt PBKDF2 should produce different hashes per call"
    # But both must verify back to the original key.
    assert verify_key_hash("ap_test123", h1)
    assert verify_key_hash("ap_test123", h2)
    assert not verify_key_hash("wrong_key", h1)