1  import random
 2  import pytest
 3  from fastapi.testclient import TestClient
 4  
 5  from restai.config import RESTAI_DEFAULT_PASSWORD
 6  from restai.main import app
 7  
 8  suffix = str(random.randint(0, 10000000))
 9  test_username = f"admin_ep_user_{suffix}"
10  test_password = "admin_ep_pass_123"
11  nonadmin_username = f"admin_ep_nonadmin_{suffix}"
12  nonadmin_password = "nonadmin_pass_123"
13  
14  ADMIN = ("admin", RESTAI_DEFAULT_PASSWORD)
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 test users for admin endpoint tests."""
25      # Create a non-admin user
26      resp = client.post(
27          "/users",
28          json={
29              "username": test_username,
30              "password": test_password,
31              "admin": False,
32              "private": False,
33          },
34          auth=ADMIN,
35      )
36      assert resp.status_code == 201
37  
38      resp = client.post(
39          "/users",
40          json={
41              "username": nonadmin_username,
42              "password": nonadmin_password,
43              "admin": False,
44              "private": False,
45          },
46          auth=ADMIN,
47      )
48      assert resp.status_code == 201
49  
50  
51  def test_impersonate():
52      """Admin can impersonate another user. Uses separate client to avoid cookie pollution."""
53      with TestClient(app) as c:
54          resp = c.post(
55              f"/auth/impersonate/{test_username}",
56              auth=ADMIN,
57          )
58          assert resp.status_code == 200
59          data = resp.json()
60          assert data["impersonating"] is True
61          assert test_username in data["message"]
62  
63  
64  def test_impersonate_nonexistent_user():
65      """Impersonating a user that doesn't exist returns 404."""
66      with TestClient(app) as c:
67          resp = c.post(
68              f"/auth/impersonate/nonexistent_user_{suffix}",
69              auth=ADMIN,
70          )
71          assert resp.status_code == 404
72  
73  
74  def test_impersonate_non_admin_rejected():
75      """Non-admin users cannot impersonate others."""
76      with TestClient(app) as c:
77          resp = c.post(
78              f"/auth/impersonate/{test_username}",
79              auth=(nonadmin_username, nonadmin_password),
80          )
81          assert resp.status_code == 403
82  
83  
84  def test_gpu_info(client):
85      """GET /settings/gpu-info returns 200 (may be empty list if no GPU)."""
86      resp = client.get("/settings/gpu-info", auth=ADMIN)
87      assert resp.status_code == 200
88      assert isinstance(resp.json(), list)
89  
90  
91  def test_cleanup(client):
92      """Remove test users."""
93      client.delete(f"/users/{test_username}", auth=ADMIN)
94      client.delete(f"/users/{nonadmin_username}", auth=ADMIN)