1  """Tests for widget chat rate limiting."""
 2  import random
 3  import pytest
 4  from fastapi.testclient import TestClient
 5  
 6  from restai.config import RESTAI_DEFAULT_PASSWORD
 7  from restai.main import app
 8  
 9  suffix = str(random.randint(0, 1000000))
10  team_name = f"wrl_team_{suffix}"
11  project_name = f"wrl_project_{suffix}"
12  team_id = None
13  project_id = None
14  widget_key = None
15  
16  
17  @pytest.fixture(scope="module")
18  def client():
19      with TestClient(app) as c:
20          yield c
21  
22  
23  def test_setup(client):
24      """Create resources for widget rate limit tests."""
25      global team_id, project_id, widget_key
26      auth = ("admin", RESTAI_DEFAULT_PASSWORD)
27      resp = client.post("/teams", json={"name": team_name}, auth=auth)
28      assert resp.status_code in (200, 201)
29      team_id = resp.json()["id"]
30  
31      resp = client.post("/projects", json={"name": project_name, "type": "block", "team_id": team_id}, auth=auth)
32      assert resp.status_code == 201
33      project_id = resp.json()["project"]
34  
35      resp = client.post(f"/projects/{project_id}/widgets", json={"name": "RL Widget", "allowed_domains": []}, auth=auth)
36      assert resp.status_code == 201
37      widget_key = resp.json()["widget_key"]
38  
39  
40  def test_widget_rate_limit_triggers(client):
41      """After 30 requests in a minute, widget gets 429."""
42      from restai.routers.widgets import _widget_requests, _widget_lock
43      with _widget_lock:
44          _widget_requests.clear()
45  
46      for i in range(30):
47          client.post("/widget/chat", json={"question": "hi"}, headers={"X-Widget-Key": widget_key})
48  
49      resp = client.post("/widget/chat", json={"question": "hi"}, headers={"X-Widget-Key": widget_key})
50      assert resp.status_code == 429
51      assert "rate limit" in resp.json()["detail"].lower()
52  
53      with _widget_lock:
54          _widget_requests.clear()
55  
56  
57  def test_widget_works_after_rate_limit_reset(client):
58      """After clearing state, widget works again."""
59      from restai.routers.widgets import _widget_requests, _widget_lock
60      with _widget_lock:
61          _widget_requests.clear()
62  
63      resp = client.post("/widget/chat", json={"question": "hi"}, headers={"X-Widget-Key": widget_key})
64      # Block project will error on chat (no LLM), but should NOT be 429
65      assert resp.status_code != 429
66  
67  
68  def test_cleanup(client):
69      """Clean up."""
70      auth = ("admin", RESTAI_DEFAULT_PASSWORD)
71      client.delete(f"/projects/{project_name}", auth=auth)
72      client.delete(f"/teams/{team_id}", auth=auth)