vinez/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-10 14:27:00 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions
headroom/tests/test_copilot_quota.py
Abhay Singh cbb775015e
fix(subscription/copilot): preserve remaining=0 for exhausted quota (#1997) 
## Description

`parse_copilot_quota` reads each category's remaining count like this
(`headroom/subscription/copilot_quota.py`):

```python
remaining = raw.get("remaining") or raw.get("quota_remaining")
```

When a Copilot category is fully consumed, the `/copilot_internal/user`
API sends
`remaining: 0`. The `or` chain treats that legitimate `0` as falsy and —
since the real
per-category payload emits `remaining`, not the `quota_remaining` alias
— collapses it to
`None`:

```python
{"entitlement": 300, "remaining": 0}   # fully spent
# raw.get("remaining") -> 0 (falsy) -> raw.get("quota_remaining") -> None -> remaining = None
```

With `remaining = None`, the derived properties break:

- `CopilotQuotaCategory.used` (needs `remaining is not None`) → `None`
instead of `entitlement`
- `used_percent`, when the API also omits `percent_remaining` for that
category → `None`

`to_dict` then emits `remaining: None, used: None, used_percent: None`,
so the dashboard
renders a **100%-exhausted** quota as `used: -` and a **0% green** gauge
— telling the user
they have full quota left when they have none.

Only the `remaining` field has this falsy-zero bug;
`entitlement`/`percent_remaining` are
already parsed with a plain `.get()`, and `overage_count`'s `or 0` is
benign because `0` is
its intended default.

Closes: no issue filed — found while auditing the subscription/quota
parsing.

## Fix

Use an explicit `is None` check, matching how the sibling fields are
parsed:

```python
remaining = raw.get("remaining")
if remaining is None:
    remaining = raw.get("quota_remaining")
```

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- `headroom/subscription/copilot_quota.py`: parse `remaining` with an
explicit `is None` check so a legitimate `0` survives (alias fallback
only when the key is truly absent).
- `tests/test_copilot_quota.py`: add
`test_fully_exhausted_remaining_zero_is_preserved` (remaining `0` →
`used == entitlement`, `used_percent == 100`).

## Testing

- [x] New regression test added (`tests/test_copilot_quota.py`)
- [x] Linting passes (`ruff check`) and formatting is clean (`ruff
format --check`) — run with the CI-pinned `ruff==0.15.17`
- [ ] Full `pytest` deferred to CI (local-OOM reason below).

```text
$ uvx ruff@0.15.17 check headroom/subscription/copilot_quota.py tests/test_copilot_quota.py
All checks passed!
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.10, headroom from this branch.
Importing `headroom` pulls in the torch/transformers stack and a full
`pytest` gets OOM-killed on this box, so I verified the parse +
`used`/`used_percent` logic with a dependency-free script and left the
full pytest to CI.
- Exact command / steps: ran a fully-exhausted category (`entitlement:
300, remaining: 0`, no alias/percent) through both the old `or`
expression and the new `is None` check, then through the
`used`/`used_percent` property logic.
- Observed result: the old path yields `remaining=None → used=None,
used_percent=None` (the misleading 0%/green); the new path preserves `0`
and reports 100%:

```text
OLD remaining: None  used=None  used_percent=None
NEW remaining: 0  used=300  used_percent=100.0
  -> OLD renders exhausted quota as unknown (0%/green); NEW shows 300/300 = 100%
OK alias fallback + normal values preserved
COPILOT QUOTA ZERO-REMAINING FIX VERIFIED
```

- Not tested: rendering the actual dashboard HTML (needs the running
app). The fix is confined to the parse function and the new test asserts
the parsed `used`/`used_percent`. Full local `pytest` deferred to CI
(OOM, per above).

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes — ran
lint + a standalone logic check; full pytest deferred to CI (local OOM,
disclosed above)
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

- One-line falsy-zero fix plus a test; no new dependencies.
- @JerrettDavis tagging you — small one, but it makes the Copilot
dashboard show a spent quota as 100% instead of a green 0%, so worth a
quick look when you have a moment.

Co-authored-by: JD Davis <mxjerrett@gmail.com>
2026-07-12 18:48:22 -04:00

348 lines
12 KiB
Python
Raw Permalink Blame History

"""Unit tests for headroom.subscription.copilot_quota."""
from __future__ import annotations
import time
import pytest
from headroom.subscription.copilot_quota import (
CopilotQuotaCategory,
CopilotQuotaSnapshot,
discover_github_token,
parse_copilot_quota,
)
# ---------------------------------------------------------------------------
# CopilotQuotaCategory helpers
# ---------------------------------------------------------------------------
class TestCopilotQuotaCategory:
def test_used_computed_from_entitlement_and_remaining(self):
cat = CopilotQuotaCategory(name="chat", entitlement=300, remaining=120)
assert cat.used == 180
def test_used_percent_computed(self):
cat = CopilotQuotaCategory(name="chat", entitlement=100, remaining=25)
assert cat.used_percent == pytest.approx(75.0)
def test_used_percent_from_percent_remaining(self):
cat = CopilotQuotaCategory(name="completions", percent_remaining=40.0)
assert cat.used_percent == pytest.approx(60.0)
def test_unlimited_used_percent_is_zero(self):
cat = CopilotQuotaCategory(name="premium_interactions", unlimited=True)
assert cat.used_percent == 0.0
def test_used_none_when_entitlement_missing(self):
cat = CopilotQuotaCategory(name="chat", remaining=50)
assert cat.used is None
def test_to_dict_keys(self):
cat = CopilotQuotaCategory(
name="chat",
entitlement=100,
remaining=60,
percent_remaining=60.0,
overage_count=2,
overage_permitted=True,
unlimited=False,
timestamp_utc="2025-01-01T00:00:00Z",
)
d = cat.to_dict()
assert d["name"] == "chat"
assert d["entitlement"] == 100
assert d["remaining"] == 60
assert d["used"] == 40
assert d["used_percent"] == pytest.approx(40.0)
assert d["overage_count"] == 2
assert d["overage_permitted"] is True
assert d["unlimited"] is False
def test_used_percent_clipped_at_zero(self):
# percent_remaining > 100 should not produce negative used_percent
cat = CopilotQuotaCategory(name="chat", percent_remaining=110.0)
assert cat.used_percent == pytest.approx(0.0)
# ---------------------------------------------------------------------------
# parse_copilot_quota
# ---------------------------------------------------------------------------
_SAMPLE_RESPONSE = {
"login": "octocat",
"copilot_plan": "individual",
"access_type_sku": "copilot_for_individuals",
"quota_reset_date_utc": "2025-02-01",
"quota_snapshots": {
"chat": {
"entitlement": 50,
"remaining": 30,
"quota_remaining": 30,
"percent_remaining": 60.0,
"overage_count": 0,
"overage_permitted": False,
"unlimited": False,
"timestamp_utc": "2025-01-15T10:00:00Z",
},
"completions": {
"entitlement": 2000,
"remaining": 1500,
"percent_remaining": 75.0,
"overage_count": 0,
"overage_permitted": True,
"unlimited": False,
"timestamp_utc": "2025-01-15T10:00:00Z",
},
"premium_interactions": {
"entitlement": 300,
"remaining": 298,
"percent_remaining": 99.3,
"overage_count": 2,
"overage_permitted": True,
"unlimited": False,
"timestamp_utc": "2025-01-15T10:00:00Z",
},
},
}
class TestParseCopilotQuota:
def test_basic_fields(self):
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
assert snap.login == "octocat"
assert snap.copilot_plan == "individual"
assert snap.access_type_sku == "copilot_for_individuals"
assert snap.quota_reset_date_utc == "2025-02-01"
def test_all_categories_parsed(self):
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
assert set(snap.categories.keys()) == {"chat", "completions", "premium_interactions"}
def test_chat_category(self):
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
chat = snap.categories["chat"]
assert chat.entitlement == 50
assert chat.remaining == 30
assert chat.percent_remaining == pytest.approx(60.0)
assert chat.unlimited is False
assert chat.overage_count == 0
def test_premium_interactions_overage(self):
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
prem = snap.categories["premium_interactions"]
assert prem.overage_count == 2
assert prem.overage_permitted is True
def test_quota_remaining_alias(self):
"""quota_remaining should be used when remaining is absent."""
data = {
"quota_snapshots": {
"chat": {
"entitlement": 100,
"quota_remaining": 75,
}
}
}
snap = parse_copilot_quota(data)
assert snap.categories["chat"].remaining == 75
def test_fully_exhausted_remaining_zero_is_preserved(self):
"""A fully-consumed category reports remaining: 0. That legitimate 0 must
survive (not become None), so used/used_percent report 100% not unknown."""
data = {
"quota_snapshots": {
"chat": {
"entitlement": 300,
"remaining": 0,
}
}
}
snap = parse_copilot_quota(data)
chat = snap.categories["chat"]
assert chat.remaining == 0
assert chat.used == 300
assert chat.used_percent == pytest.approx(100.0)
def test_unlimited_category(self):
data = {"quota_snapshots": {"completions": {"unlimited": True}}}
snap = parse_copilot_quota(data)
assert snap.categories["completions"].unlimited is True
def test_empty_quota_snapshots(self):
snap = parse_copilot_quota({"login": "ghost"})
assert snap.login == "ghost"
assert snap.categories == {}
def test_quota_reset_date_fallback(self):
data = {"quota_reset_date": "2025-03-01"}
snap = parse_copilot_quota(data)
assert snap.quota_reset_date_utc == "2025-03-01"
def test_fetched_at_is_recent(self):
before = time.time()
snap = parse_copilot_quota({})
after = time.time()
assert before <= snap.fetched_at <= after
def test_to_dict_structure(self):
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
d = snap.to_dict()
assert "login" in d
assert "categories" in d
assert "chat" in d["categories"]
assert "used_percent" in d["categories"]["chat"]
def test_missing_categories_skipped(self):
data = {
"quota_snapshots": {
"chat": {"remaining": 10},
# completions and premium_interactions absent
}
}
snap = parse_copilot_quota(data)
assert "chat" in snap.categories
assert "completions" not in snap.categories
assert "premium_interactions" not in snap.categories
def test_free_plan(self):
data = {"copilot_plan": "free", "quota_snapshots": {}}
snap = parse_copilot_quota(data)
assert snap.copilot_plan == "free"
# ---------------------------------------------------------------------------
# discover_github_token
# ---------------------------------------------------------------------------
class TestDiscoverGithubToken:
def test_returns_none_when_no_env_vars(self, monkeypatch):
for var in [
"GITHUB_COPILOT_GITHUB_TOKEN",
"GITHUB_TOKEN",
"COPILOT_GITHUB_TOKEN",
"GITHUB_COPILOT_API_TOKEN",
]:
monkeypatch.delenv(var, raising=False)
assert discover_github_token() is None
def test_picks_up_github_token(self, monkeypatch):
for var in [
"GITHUB_COPILOT_GITHUB_TOKEN",
"GITHUB_TOKEN",
"COPILOT_GITHUB_TOKEN",
"GITHUB_COPILOT_API_TOKEN",
]:
monkeypatch.delenv(var, raising=False)
monkeypatch.setenv("GITHUB_TOKEN", "ghp_testtoken123")
assert discover_github_token() == "ghp_testtoken123"
def test_prefers_copilot_specific_token(self, monkeypatch):
for var in [
"GITHUB_COPILOT_GITHUB_TOKEN",
"GITHUB_TOKEN",
"COPILOT_GITHUB_TOKEN",
"GITHUB_COPILOT_API_TOKEN",
]:
monkeypatch.delenv(var, raising=False)
monkeypatch.setenv("GITHUB_COPILOT_GITHUB_TOKEN", "ghp_copilot_specific")
monkeypatch.setenv("GITHUB_TOKEN", "ghp_generic")
assert discover_github_token() == "ghp_copilot_specific"
def test_falls_through_to_next_env_var(self, monkeypatch):
for var in [
"GITHUB_COPILOT_GITHUB_TOKEN",
"GITHUB_TOKEN",
"COPILOT_GITHUB_TOKEN",
"GITHUB_COPILOT_API_TOKEN",
]:
monkeypatch.delenv(var, raising=False)
monkeypatch.setenv("COPILOT_GITHUB_TOKEN", "ghp_copilot")
assert discover_github_token() == "ghp_copilot"
def test_ignores_empty_strings(self, monkeypatch):
for var in [
"GITHUB_COPILOT_GITHUB_TOKEN",
"GITHUB_TOKEN",
"COPILOT_GITHUB_TOKEN",
"GITHUB_COPILOT_API_TOKEN",
]:
monkeypatch.delenv(var, raising=False)
monkeypatch.setenv("GITHUB_COPILOT_GITHUB_TOKEN", "")
monkeypatch.setenv("GITHUB_TOKEN", "ghp_valid")
assert discover_github_token() == "ghp_valid"
# ---------------------------------------------------------------------------
# CopilotQuotaSnapshot.to_dict
# ---------------------------------------------------------------------------
class TestCopilotQuotaSnapshot:
def test_to_dict_complete(self):
snap = CopilotQuotaSnapshot(
login="user1",
copilot_plan="business",
access_type_sku="copilot_enterprise",
quota_reset_date_utc="2025-02-01",
)
snap.categories["chat"] = CopilotQuotaCategory(name="chat", entitlement=50, remaining=25)
d = snap.to_dict()
assert d["login"] == "user1"
assert d["copilot_plan"] == "business"
assert "chat" in d["categories"]
assert d["categories"]["chat"]["entitlement"] == 50
# ---------------------------------------------------------------------------
# Poll-loop task-leak regression
# ---------------------------------------------------------------------------
class TestCopilotQuotaPollLoopLeak:
@pytest.mark.asyncio
async def test_poll_loop_does_not_leak_event_wait_tasks(self, monkeypatch):
"""Regression for the ``asyncio.shield(event.wait())`` pattern.
Matches the equivalent guard in ``tests/test_subscription_tracker.py``:
every poll interval the loop previously leaked one Event.wait
waiter because ``asyncio.shield`` prevented ``wait_for`` from
cancelling the inner wait on timeout.
"""
import asyncio
from headroom.subscription.copilot_quota import _CopilotQuotaTracker
# No token configured → _maybe_poll returns immediately each cycle.
for var in ("GITHUB_COPILOT_GITHUB_TOKEN", "GITHUB_TOKEN"):
monkeypatch.delenv(var, raising=False)
tracker = _CopilotQuotaTracker(poll_interval_s=0.05)
def _count_event_wait() -> int:
return sum(
1
for t in asyncio.all_tasks()
if (t.get_coro().__qualname__ if t.get_coro() else "") == "Event.wait"
)
baseline = _count_event_wait()
await tracker.start()
try:
await asyncio.sleep(0.3) # ~6 poll cycles
peak = _count_event_wait()
finally:
await tracker.stop()
await asyncio.sleep(0.05)
residual = _count_event_wait()
assert peak - baseline <= 1, (
f"CopilotQuotaTracker leaked Event.wait: baseline={baseline} peak={peak}"
)
assert residual <= baseline, (
f"CopilotQuotaTracker left residual Event.wait: baseline={baseline} residual={residual}"
)
Reference in a new issue View git blame Copy permalink