| 1 | """Tests for token accounting policy: estimates, billing direction, sums."""
|
| 2 |
|
| 3 | from __future__ import annotations
|
| 4 |
|
| 5 | import sys
|
| 6 | from pathlib import Path
|
| 7 |
|
| 8 | sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts"))
|
| 9 |
|
| 10 | import unittest
|
| 11 |
|
| 12 | from oakbench.tokens import (
|
| 13 | COST_WEIGHTS_NOTE,
|
| 14 | cost_weighted_total,
|
| 15 | estimate_tokens,
|
| 16 | int_or_zero,
|
| 17 | interaction_token_fields,
|
| 18 | summed_token_fields,
|
| 19 | )
|
| 20 | from oakbench.environment import command_display
|
| 21 |
|
| 22 |
|
| 23 | class EstimateTest(unittest.TestCase):
|
| 24 | def test_empty_is_zero(self) -> None:
|
| 25 | self.assertEqual(estimate_tokens(""), 0)
|
| 26 |
|
| 27 | def test_rounds_up_and_floors_at_one(self) -> None:
|
| 28 | self.assertEqual(estimate_tokens("a"), 1)
|
| 29 | self.assertEqual(estimate_tokens("abcd"), 1)
|
| 30 | self.assertEqual(estimate_tokens("abcde"), 2)
|
| 31 |
|
| 32 |
|
| 33 | class CostWeightTest(unittest.TestCase):
|
| 34 | def test_billing_direction(self) -> None:
|
| 35 | # Emitted (model output) is 5x ingested (model input); cache reads 0.1x.
|
| 36 | self.assertEqual(cost_weighted_total(10, 100), 150.0)
|
| 37 | self.assertEqual(cost_weighted_total(0, 0, 1000), 100.0)
|
| 38 |
|
| 39 | def test_weights_note_is_explicit(self) -> None:
|
| 40 | self.assertIn("emitted=5.0", COST_WEIGHTS_NOTE)
|
| 41 | self.assertIn("ingested=1.0", COST_WEIGHTS_NOTE)
|
| 42 |
|
| 43 |
|
| 44 | class CommandDisplayPricingTest(unittest.TestCase):
|
| 45 | def test_argv_display_prices_typed_executable_name(self) -> None:
|
| 46 | display = command_display(["/private/tmp/oak-subject/target/debug/oak", "pull"])
|
| 47 |
|
| 48 | self.assertEqual(display, "oak pull")
|
| 49 | fields = interaction_token_fields(display, "", "", 0, 0, False, False)
|
| 50 | self.assertEqual(fields["estimated_tokens_agent_emitted"], estimate_tokens("oak pull"))
|
| 51 |
|
| 52 | def test_only_executable_path_is_shortened(self) -> None:
|
| 53 | self.assertEqual(
|
| 54 | command_display(["/usr/bin/git", "clone", "/private/tmp/source-repo", "dst"]),
|
| 55 | "git clone /private/tmp/source-repo dst",
|
| 56 | )
|
| 57 |
|
| 58 | def test_shell_and_synthetic_labels_stay_verbatim(self) -> None:
|
| 59 | self.assertEqual(command_display(["/bin/zsh", "-lc", "/private/tmp/oak pull"]), "/private/tmp/oak pull")
|
| 60 | self.assertEqual(command_display(["/bin/zsh", "-c", "/private/tmp/oak pull"]), "/private/tmp/oak pull")
|
| 61 | self.assertEqual(
|
| 62 | command_display("codex_exec_json network: /private/tmp/oak pull"),
|
| 63 | "codex_exec_json network: /private/tmp/oak pull",
|
| 64 | )
|
| 65 |
|
| 66 |
|
| 67 | class InteractionFieldsTest(unittest.TestCase):
|
| 68 | def test_directions_and_compat_names(self) -> None:
|
| 69 | fields = interaction_token_fields("git status", "out" * 10, "", 30, 0, False, False)
|
| 70 | # The command the agent types is model-emitted; stdout is ingested.
|
| 71 | self.assertEqual(fields["estimated_tokens_agent_emitted"], estimate_tokens("git status"))
|
| 72 | self.assertEqual(fields["estimated_tokens_agent_ingested"], estimate_tokens("out" * 10))
|
| 73 | # Deprecated CLI-perspective names stay equal for compatibility.
|
| 74 | self.assertEqual(fields["estimated_tokens_input"], fields["estimated_tokens_agent_emitted"])
|
| 75 | self.assertEqual(fields["estimated_tokens_output"], fields["estimated_tokens_agent_ingested"])
|
| 76 | self.assertEqual(
|
| 77 | fields["estimated_cost_weighted_tokens"],
|
| 78 | cost_weighted_total(fields["estimated_tokens_agent_emitted"], fields["estimated_tokens_agent_ingested"]),
|
| 79 | )
|
| 80 | self.assertEqual(fields["raw_output_bytes"], 30)
|
| 81 | self.assertIn("direction_note", fields["token_estimate"])
|
| 82 |
|
| 83 | def test_truncation_flags_propagate(self) -> None:
|
| 84 | fields = interaction_token_fields("x", "y", "z", 1, 1, True, False)
|
| 85 | self.assertTrue(fields["output_truncated"])
|
| 86 | self.assertTrue(fields["stdout_truncated"])
|
| 87 | self.assertFalse(fields["stderr_truncated"])
|
| 88 |
|
| 89 |
|
| 90 | class SummedFieldsTest(unittest.TestCase):
|
| 91 | def test_int_or_zero_handles_nulls_and_bad_values(self) -> None:
|
| 92 | self.assertEqual(int_or_zero(None), 0)
|
| 93 | self.assertEqual(int_or_zero("7"), 7)
|
| 94 | self.assertEqual(int_or_zero(True), 0)
|
| 95 | self.assertEqual(int_or_zero(-3), 0)
|
| 96 |
|
| 97 | def test_sum_equals_parts(self) -> None:
|
| 98 | parts = [
|
| 99 | interaction_token_fields("git add .", "", "", 0, 0, False, False),
|
| 100 | interaction_token_fields("git commit -m initial", "done", "", 4, 0, False, False),
|
| 101 | ]
|
| 102 | total = summed_token_fields(parts, "test_sum")
|
| 103 | self.assertEqual(
|
| 104 | total["estimated_tokens_total"],
|
| 105 | sum(p["estimated_tokens_total"] for p in parts),
|
| 106 | )
|
| 107 | self.assertEqual(
|
| 108 | total["estimated_cost_weighted_tokens"],
|
| 109 | cost_weighted_total(
|
| 110 | sum(p["estimated_tokens_input"] for p in parts),
|
| 111 | sum(p["estimated_tokens_output"] for p in parts),
|
| 112 | ),
|
| 113 | )
|
| 114 | self.assertEqual(total["raw_output_bytes"], 4)
|
| 115 | self.assertEqual(total["token_estimate"]["source"], "test_sum")
|
| 116 |
|
| 117 | def test_sum_treats_null_unmeasured_fields_as_zero_contribution(self) -> None:
|
| 118 | total = summed_token_fields(
|
| 119 | [
|
| 120 | {
|
| 121 | "estimated_tokens_input": None,
|
| 122 | "estimated_tokens_output": None,
|
| 123 | "stdout_bytes": None,
|
| 124 | "stderr_bytes": None,
|
| 125 | "estimated_tokens_envelope_emitted": None,
|
| 126 | "estimated_tokens_envelope_ingested": None,
|
| 127 | "token_estimate": {
|
| 128 | "stdout_admitted_chars": None,
|
| 129 | "stderr_admitted_chars": None,
|
| 130 | },
|
| 131 | },
|
| 132 | interaction_token_fields("oak status", "ok", "", 2, 0, False, False),
|
| 133 | ],
|
| 134 | "test_null_sum",
|
| 135 | command_chars=10,
|
| 136 | )
|
| 137 | self.assertEqual(total["estimated_tokens_total"], estimate_tokens("oak status") + estimate_tokens("ok"))
|
| 138 | self.assertEqual(total["raw_output_bytes"], 2)
|
| 139 | self.assertEqual(total["token_estimate"]["stdout_admitted_chars"], len("ok"))
|
| 140 |
|
| 141 |
|
| 142 | if __name__ == "__main__":
|
| 143 | unittest.main()
|