openclaw/skills/model-usage/scripts/test_model_usage.py

#!/usr/bin/env python3
"""
Tests for model_usage helpers.
"""

import argparse
from datetime import date, timedelta
from unittest import TestCase, main

from model_usage import (
    aggregate_costs,
    coerce_finite_cost,
    filter_by_days,
    latest_day_cost,
    pick_current_model,
    positive_int,
)


class TestModelUsage(TestCase):
    def test_positive_int_accepts_valid_numbers(self):
        self.assertEqual(positive_int("1"), 1)
        self.assertEqual(positive_int("7"), 7)

    def test_positive_int_rejects_zero_and_negative(self):
        with self.assertRaises(argparse.ArgumentTypeError):
            positive_int("0")
        with self.assertRaises(argparse.ArgumentTypeError):
            positive_int("-3")

    def test_filter_by_days_keeps_recent_entries(self):
        today = date.today()
        entries = [
            {"date": (today - timedelta(days=5)).strftime("%Y-%m-%d"), "modelBreakdowns": []},
            {"date": (today - timedelta(days=1)).strftime("%Y-%m-%d"), "modelBreakdowns": []},
            {"date": today.strftime("%Y-%m-%d"), "modelBreakdowns": []},
        ]

        filtered = filter_by_days(entries, 2)

        self.assertEqual(len(filtered), 2)
        self.assertEqual(filtered[0]["date"], (today - timedelta(days=1)).strftime("%Y-%m-%d"))
        self.assertEqual(filtered[1]["date"], today.strftime("%Y-%m-%d"))

    def test_coerce_finite_cost_accepts_numbers_and_numeric_strings(self):
        self.assertEqual(coerce_finite_cost(2), 2.0)
        self.assertEqual(coerce_finite_cost(1.75), 1.75)
        self.assertEqual(coerce_finite_cost("1.75"), 1.75)
        self.assertEqual(coerce_finite_cost("  2.5 "), 2.5)

    def test_coerce_finite_cost_rejects_booleans(self):
        # bool is a subclass of int in Python, but is never a valid cost.
        self.assertIsNone(coerce_finite_cost(True))
        self.assertIsNone(coerce_finite_cost(False))

    def test_coerce_finite_cost_rejects_non_finite(self):
        self.assertIsNone(coerce_finite_cost(float("nan")))
        self.assertIsNone(coerce_finite_cost(float("inf")))
        self.assertIsNone(coerce_finite_cost(float("-inf")))
        self.assertIsNone(coerce_finite_cost("NaN"))
        self.assertIsNone(coerce_finite_cost("Infinity"))

    def test_coerce_finite_cost_rejects_unusable_values(self):
        self.assertIsNone(coerce_finite_cost("not-a-number"))
        self.assertIsNone(coerce_finite_cost(""))
        self.assertIsNone(coerce_finite_cost(None))
        self.assertIsNone(coerce_finite_cost({}))

    def test_aggregate_costs_includes_numeric_strings(self):
        entries = [
            {
                "date": "2026-05-25",
                "modelBreakdowns": [
                    {"modelName": "claude-sonnet-4-6", "cost": 1.50},
                    {"modelName": "claude-sonnet-4-6", "cost": "1.75"},
                ],
            }
        ]
        self.assertEqual(aggregate_costs(entries), {"claude-sonnet-4-6": 3.25})

    def test_aggregate_costs_ignores_bool_and_non_finite(self):
        entries = [
            {
                "date": "2026-05-25",
                "modelBreakdowns": [
                    {"modelName": "claude-sonnet-4-6", "cost": 1.50},
                    {"modelName": "claude-sonnet-4-6", "cost": "1.75"},
                    {"modelName": "claude-sonnet-4-6", "cost": True},
                    {"modelName": "claude-sonnet-4-6", "cost": float("nan")},
                    {"modelName": "claude-sonnet-4-6", "cost": float("inf")},
                ],
            }
        ]
        totals = aggregate_costs(entries)
        # NaN/Infinity must not poison the total; bool must not add 1.0.
        self.assertEqual(totals, {"claude-sonnet-4-6": 3.25})

    def test_pick_current_model_scores_numeric_string_costs(self):
        # model-b's cost is a numeric string; it must still win on highest cost.
        entries = [
            {
                "date": "2026-05-25",
                "modelBreakdowns": [
                    {"modelName": "model-a", "cost": 1.0},
                    {"modelName": "model-b", "cost": "5.0"},
                ],
            }
        ]
        model, day = pick_current_model(entries)
        self.assertEqual(model, "model-b")
        self.assertEqual(day, "2026-05-25")

    def test_pick_current_model_ignores_bool_and_non_finite(self):
        # Only model-a has a usable cost; bool and NaN must not be scored.
        entries = [
            {
                "date": "2026-05-25",
                "modelBreakdowns": [
                    {"modelName": "model-a", "cost": 2.0},
                    {"modelName": "model-b", "cost": True},
                    {"modelName": "model-c", "cost": float("nan")},
                ],
            }
        ]
        model, _day = pick_current_model(entries)
        self.assertEqual(model, "model-a")

    def test_latest_day_cost_accepts_numeric_string(self):
        entries = [
            {
                "date": "2026-05-25",
                "modelBreakdowns": [{"modelName": "model-a", "cost": "2.50"}],
            }
        ]
        day, cost = latest_day_cost(entries, "model-a")
        self.assertEqual(day, "2026-05-25")
        self.assertEqual(cost, 2.50)

    def test_latest_day_cost_rejects_non_finite(self):
        entries = [
            {
                "date": "2026-05-25",
                "modelBreakdowns": [{"modelName": "model-a", "cost": float("inf")}],
            }
        ]
        day, cost = latest_day_cost(entries, "model-a")
        self.assertEqual(day, "2026-05-25")
        self.assertIsNone(cost)


if __name__ == "__main__":
    main()