1  """Tests for per-project budget checks (rate_limit is already covered by test_rate_limit.py)."""
 2  
 3  import random
 4  import pytest
 5  from fastapi.testclient import TestClient
 6  
 7  from restai.config import RESTAI_DEFAULT_PASSWORD
 8  from restai.main import app
 9  
10  ADMIN = ("admin", RESTAI_DEFAULT_PASSWORD)
11  
12  suffix = str(random.randint(0, 999999))
13  team_name = f"budget_team_{suffix}"
14  llm_name = f"budget_llm_{suffix}"
15  project_name = f"budget-proj-{suffix}"
16  
17  team_id = None
18  project_id = None
19  
20  
21  @pytest.fixture(scope="module")
22  def client():
23      with TestClient(app) as c:
24          yield c
25  
26  
27  def test_setup(client):
28      """Create team, LLM, and block project."""
29      global team_id, project_id
30  
31      r = client.post(
32          "/llms",
33          json={
34              "name": llm_name,
35              "class_name": "OpenAI",
36              "options": {"model": "gpt-test", "api_key": "sk-fake"},
37              "privacy": "public",
38          },
39          auth=ADMIN,
40      )
41      assert r.status_code in (200, 201)
42  
43      r = client.post(
44          "/teams",
45          json={"name": team_name},
46          auth=ADMIN,
47      )
48      assert r.status_code == 201
49      team_id = r.json()["id"]
50  
51      r = client.post(
52          "/projects",
53          json={"name": project_name, "type": "block", "team_id": team_id},
54          auth=ADMIN,
55      )
56      assert r.status_code == 201
57      project_id = r.json()["project"]
58  
59  
60  def test_project_options_rate_limit_persists(client):
61      """Verify rate_limit option is saved and returned in project details."""
62      r = client.patch(
63          f"/projects/{project_id}",
64          json={"options": {"rate_limit": 50}},
65          auth=ADMIN,
66      )
67      assert r.status_code == 200
68  
69      r = client.get(f"/projects/{project_id}", auth=ADMIN)
70      assert r.status_code == 200
71      assert r.json()["options"]["rate_limit"] == 50
72  
73  
74  def test_project_options_rate_limit_nullable(client):
75      """Verify rate_limit can be set to null (unlimited)."""
76      r = client.patch(
77          f"/projects/{project_id}",
78          json={"options": {"rate_limit": None}},
79          auth=ADMIN,
80      )
81      assert r.status_code == 200
82  
83      r = client.get(f"/projects/{project_id}", auth=ADMIN)
84      assert r.status_code == 200
85      assert r.json()["options"]["rate_limit"] is None
86  
87  
88  def test_cleanup(client):
89      """Clean up resources."""
90      client.delete(f"/projects/{project_id}", auth=ADMIN)
91      client.delete(f"/llms/{llm_name}", auth=ADMIN)
92      client.delete(f"/teams/{team_id}", auth=ADMIN)