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  
  9  @pytest.fixture(autouse=True)
 10  def clear_rate_limiter():
 11      """Clear DB-backed login rate limiter before each test."""
 12      from restai.database import DBWrapper
 13      from restai.models.databasemodels import LoginAttemptDatabase
 14      db = DBWrapper()
 15      try:
 16          db.db.query(LoginAttemptDatabase).delete()
 17          db.db.commit()
 18      except Exception:
 19          db.db.rollback()
 20      finally:
 21          db.db.close()
 22  
 23  
 24  test_username = "test_auth_user_" + str(random.randint(0, 1000000))
 25  test_password = "auth_test_pass"
 26  
 27  
 28  def test_setup_user():
 29      with TestClient(app) as client:
 30          response = client.post(
 31              "/users",
 32              json={
 33                  "username": test_username,
 34                  "password": test_password,
 35                  "admin": False,
 36                  "private": False,
 37              },
 38              auth=("admin", RESTAI_DEFAULT_PASSWORD),
 39          )
 40          assert response.status_code == 201
 41          assert response.json()["username"] == test_username
 42  
 43  
 44  def test_login():
 45      with TestClient(app) as client:
 46          response = client.post(
 47              "/auth/login",
 48              auth=(test_username, test_password),
 49          )
 50          assert response.status_code == 200
 51          assert response.json()["message"] == "Logged in successfully."
 52          assert "restai_token" in response.cookies
 53  
 54  
 55  def test_whoami_with_cookie():
 56      with TestClient(app) as client:
 57          login_resp = client.post(
 58              "/auth/login",
 59              auth=(test_username, test_password),
 60          )
 61          cookie = login_resp.cookies.get("restai_token")
 62          assert cookie is not None
 63  
 64          response = client.get(
 65              "/auth/whoami",
 66              cookies={"restai_token": cookie},
 67          )
 68          assert response.status_code == 200
 69          assert response.json()["username"] == test_username
 70  
 71  
 72  def test_whoami_with_basic_auth():
 73      with TestClient(app) as client:
 74          response = client.get(
 75              "/auth/whoami",
 76              auth=(test_username, test_password),
 77          )
 78          assert response.status_code == 200
 79          assert response.json()["username"] == test_username
 80  
 81  
 82  def test_whoami_admin():
 83      with TestClient(app) as client:
 84          response = client.get(
 85              "/auth/whoami",
 86              auth=("admin", RESTAI_DEFAULT_PASSWORD),
 87          )
 88          assert response.status_code == 200
 89          data = response.json()
 90          assert data["username"] == "admin"
 91          assert data["is_admin"] is True
 92  
 93  
 94  def test_login_wrong_password():
 95      with TestClient(app) as client:
 96          response = client.post(
 97              "/auth/login",
 98              auth=(test_username, "wrong_password"),
 99          )
100          assert response.status_code == 401
101  
102  
103  def test_whoami_no_auth():
104      with TestClient(app) as client:
105          response = client.get("/auth/whoami")
106          assert response.status_code == 401
107  
108  
109  def test_logout_and_cleanup():
110      with TestClient(app) as client:
111          login_resp = client.post(
112              "/auth/login",
113              auth=(test_username, test_password),
114          )
115          cookie = login_resp.cookies.get("restai_token")
116  
117          response = client.post(
118              "/auth/logout",
119              cookies={"restai_token": cookie},
120          )
121          assert response.status_code == 200
122          assert response.json()["message"] == "Logged out successfully."
123  
124          # Clean up test user
125          response = client.delete(
126              f"/users/{test_username}",
127              auth=("admin", RESTAI_DEFAULT_PASSWORD),
128          )
129          assert response.status_code == 200