adaptive-pricing/projects/coulomb-pricing/tests/test_customer_tuning.py

from __future__ import annotations

from decimal import Decimal
from pathlib import Path

from adaptive_pricing_core.customer_tuning import CustomerTuningRequest, solve_customer_tuning
from observatory.boundary import build_boundary_policy
from observatory.economics import build_snapshot
from observatory.load import (
    load_ltv_scenarios,
    load_membership,
    load_monthly_ledger,
    load_payment_records,
    load_pricing_models,
    load_product,
)
from observatory.ltv import _configuration, _ltv_policy, _profile, _usage_unit_cost
from observatory.tuning import build_customer_tuning_pilot
from observatory.usage import load_usage_records

DATA_DIR = Path(__file__).resolve().parent.parent / "data"


def _scenario_inputs(profile_id: str = "small-team"):
    product = load_product(DATA_DIR)
    models = load_pricing_models(DATA_DIR)
    members = load_membership(DATA_DIR)
    payments = load_payment_records(DATA_DIR)
    ledger = load_monthly_ledger(DATA_DIR)
    snapshot = build_snapshot("2026-06", product, models, members, payments, ledger)
    usage_records = load_usage_records(DATA_DIR)
    scenario_catalog = load_ltv_scenarios(DATA_DIR)
    profile = _profile(next(item for item in scenario_catalog["profiles"] if item["id"] == profile_id))
    observed_usage_unit_cost = _usage_unit_cost(usage_records, snapshot.period)
    return snapshot, models, usage_records, scenario_catalog, profile, observed_usage_unit_cost


def test_lower_usage_price_request_can_stay_seller_safe() -> None:
    (
        snapshot,
        models,
        usage_records,
        scenario_catalog,
        profile,
        observed_usage_unit_cost,
    ) = _scenario_inputs()
    model = next(item for item in models if item.id == "membership-plus-overage")
    outcome = solve_customer_tuning(
        _configuration(model, profile, snapshot, observed_usage_unit_cost),
        [
            _configuration(candidate, profile, snapshot, observed_usage_unit_cost)
            for candidate in models
            if candidate.status in ("active", "candidate")
        ],
        profile,
        build_boundary_policy(snapshot),
        _ltv_policy(scenario_catalog),
        CustomerTuningRequest(
            included_units=Decimal("50000"),
            contract_duration_months=3,
            preference="lower_usage_price",
            approval_mode="self_serve_only",
        ),
    )

    assert outcome.decision == "accepted"
    assert outcome.reference_model_id == "flat-899-eur-monthly"
    assert outcome.passes_required_improvement is True
    assert outcome.solved_usage_unit_price < Decimal("0.002")
    assert "lower_included_usage" in outcome.tradeoffs
    assert "longer_contract_duration" in outcome.tradeoffs


def test_high_included_request_is_rejected_for_self_serve() -> None:
    (
        snapshot,
        models,
        _usage_records,
        scenario_catalog,
        profile,
        observed_usage_unit_cost,
    ) = _scenario_inputs()
    model = next(item for item in models if item.id == "membership-plus-overage")
    outcome = solve_customer_tuning(
        _configuration(model, profile, snapshot, observed_usage_unit_cost),
        [
            _configuration(candidate, profile, snapshot, observed_usage_unit_cost)
            for candidate in models
            if candidate.status in ("active", "candidate")
        ],
        profile,
        build_boundary_policy(snapshot),
        _ltv_policy(scenario_catalog),
        CustomerTuningRequest(
            included_units=Decimal("150000"),
            contract_duration_months=3,
            preference="lower_usage_price",
            approval_mode="self_serve_only",
        ),
    )

    assert outcome.decision == "rejected"
    assert outcome.passes_required_improvement is True
    assert any(
        constraint.id == "discount-exposure-limit"
        for constraint in outcome.binding_constraints
    )


def test_customer_tuning_pilot_surfaces_accepted_and_rejected_requests() -> None:
    snapshot, models, usage_records, scenario_catalog, _profile_data, _usage_unit_cost_value = _scenario_inputs()
    pilot = build_customer_tuning_pilot(
        snapshot,
        models,
        usage_records,
        scenario_catalog,
        {
            "requests": [
                {
                    "id": "accepted",
                    "name": "Accepted",
                    "profile_id": "small-team",
                    "model_id": "membership-plus-overage",
                    "preference": "lower_usage_price",
                    "approval_mode": "self_serve_only",
                    "selected_tunables": {
                        "included_tokens": "50000",
                        "contract_duration_months": 3,
                    },
                },
                {
                    "id": "rejected",
                    "name": "Rejected",
                    "profile_id": "small-team",
                    "model_id": "membership-plus-overage",
                    "preference": "lower_usage_price",
                    "approval_mode": "self_serve_only",
                    "selected_tunables": {
                        "included_tokens": "150000",
                        "contract_duration_months": 3,
                    },
                },
            ]
        },
    )

    assert pilot["request_count"] == 2
    assert pilot["accepted_request_ids"] == ["accepted"]
    assert {item["result"]["decision"] for item in pilot["requests"]} == {"accepted", "rejected"}