File size: 4,902 Bytes
84013a9
 
 
 
af423c3
 
 
 
 
 
 
 
 
 
 
84013a9
af423c3
84013a9
 
 
 
 
 
 
af423c3
84013a9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
af423c3
84013a9
 
 
 
 
 
f8e7d59
84013a9
 
 
 
 
 
 
 
 
f8e7d59
84013a9
 
 
 
 
 
 
f8e7d59
84013a9
 
 
 
 
f8e7d59
84013a9
 
 
 
f8e7d59
84013a9
 
f8e7d59
84013a9
 
 
 
 
 
f8e7d59
84013a9
 
 
 
6fbc8ca
 
 
 
f8e7d59
6fbc8ca
 
 
 
 
 
 
 
 
f8e7d59
6fbc8ca
 
 
 
 
 
 
 
f8e7d59
6fbc8ca
 
 
 
 
 
 
 
f8e7d59
6fbc8ca
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
"""
Route-level tests for the /admin/keys endpoints against the real,
Postgres-backed UsageTracker -- these routes are normally exercised through
the module-level `tracker` singleton (conftest.py replaces it globally with
MockTracker for every other test), so `tracker` is patched on
`app.core.usage_tracker` and `ADMIN_API_KEY` on `app.api.routes_admin` for
the duration of this module.

`tracker` is patched on the *defining* module, not on routes_admin, because
routes_admin no longer holds its own binding: it reaches the singleton
through `usage_tracker.tracker` so that init_tracker()'s rebinding is
visible to it. That is also what makes one patch here sufficient where five
would otherwise be needed.

This is the main regression coverage for the api_keys-on-Postgres
migration: it exercises the exact raw SQL routes_admin.py runs against
`usage_tracker.tracker._get_pg_conn()`.
"""
import hashlib
import hmac
import os

import pytest

from app.core import usage_tracker
from app.core.usage_tracker import UsageTracker
import app.api.routes_admin as routes_admin

TEST_ADMIN_KEY = "test-admin-key-for-routes-admin-tests"
TEST_PEPPER = os.environ["ARF_KEY_PEPPER"]  # set in conftest.py before app import


def _key_id(raw_key: str) -> str:
    """Reproduce UsageTracker._lookup_hash without a tracker instance, so
    tests can locate the row created for a given raw key deterministically."""
    return hmac.new(TEST_PEPPER.encode(), raw_key.encode(), hashlib.sha256).hexdigest()


@pytest.fixture(autouse=True)
def real_tracker(monkeypatch):
    real = UsageTracker(db_path=":memory:")
    monkeypatch.setattr(usage_tracker, "tracker", real)
    monkeypatch.setattr(routes_admin, "ADMIN_API_KEY", TEST_ADMIN_KEY)
    yield real


def test_create_list_update_deactivate_key(client):
    create_resp = client.post(
        "/api/v1/admin/keys",
        params={"admin_key": TEST_ADMIN_KEY},
        json={"tier": "free", "org_name": "Test Org"},
    )
    assert create_resp.status_code == 200
    body = create_resp.json()
    api_key = body["api_key"]
    assert body["tier"] == "free"
    key_id = _key_id(api_key)

    list_resp = client.get("/api/v1/admin/keys", params={"admin_key": TEST_ADMIN_KEY})
    assert list_resp.status_code == 200
    keys_by_id = {row["key_id"]: row for row in list_resp.json()["keys"]}
    assert key_id in keys_by_id
    assert keys_by_id[key_id]["tier"] == "free"
    assert keys_by_id[key_id]["is_active"] is True

    patch_resp = client.patch(
        f"/api/v1/admin/keys/{key_id}/tier",
        params={"admin_key": TEST_ADMIN_KEY},
        json={"tier": "pro"},
    )
    assert patch_resp.status_code == 200

    list_resp2 = client.get("/api/v1/admin/keys", params={"admin_key": TEST_ADMIN_KEY})
    assert list_resp2.json()["keys"][0]  # non-empty, sanity check
    keys_by_id2 = {row["key_id"]: row for row in list_resp2.json()["keys"]}
    assert keys_by_id2[key_id]["tier"] == "pro"

    delete_resp = client.delete(f"/api/v1/admin/keys/{key_id}", params={"admin_key": TEST_ADMIN_KEY})
    assert delete_resp.status_code == 200

    list_resp3 = client.get("/api/v1/admin/keys", params={"admin_key": TEST_ADMIN_KEY})
    keys_by_id3 = {row["key_id"]: row for row in list_resp3.json()["keys"]}
    assert keys_by_id3[key_id]["is_active"] is False


def test_update_nonexistent_key_returns_404(client):
    resp = client.patch(
        "/api/v1/admin/keys/does-not-exist/tier",
        params={"admin_key": TEST_ADMIN_KEY},
        json={"tier": "pro"},
    )
    assert resp.status_code == 404


def test_rotate_key_deactivates_old_and_creates_new_on_same_tenant(client):
    create_resp = client.post(
        "/api/v1/admin/keys",
        params={"admin_key": TEST_ADMIN_KEY},
        json={"tier": "pro", "org_name": "Rotate Test Org"},
    )
    assert create_resp.status_code == 200
    old_body = create_resp.json()
    old_key_id = _key_id(old_body["api_key"])
    tenant_id = old_body["tenant_id"]

    rotate_resp = client.post(
        f"/api/v1/admin/keys/{old_key_id}/rotate", params={"admin_key": TEST_ADMIN_KEY})
    assert rotate_resp.status_code == 200
    rotated = rotate_resp.json()
    assert rotated["tenant_id"] == tenant_id
    assert rotated["tier"] == "pro"
    assert rotated["deactivated_key_id"] == old_key_id
    new_key_id = _key_id(rotated["api_key"])
    assert new_key_id != old_key_id

    list_resp = client.get("/api/v1/admin/keys", params={"admin_key": TEST_ADMIN_KEY})
    keys_by_id = {row["key_id"]: row for row in list_resp.json()["keys"]}
    assert keys_by_id[old_key_id]["is_active"] is False
    assert keys_by_id[new_key_id]["is_active"] is True
    assert keys_by_id[new_key_id]["tier"] == "pro"


def test_rotate_nonexistent_key_returns_404(client):
    resp = client.post(
        "/api/v1/admin/keys/does-not-exist/rotate", params={"admin_key": TEST_ADMIN_KEY})
    assert resp.status_code == 404