Gan, Jimmy
b981c06d59
- Fix unit test imports: add env setup in conftest.py before module imports - Add 24 new auth router tests (RBAC, preferences, password validation) - Add 16 new tests for litellm and chat_summary routers - Apply black formatting and ruff linting across codebase - Add pre-commit hooks configuration (black, ruff, file checks) - Increase CI coverage threshold from 40% to 50% Test Results: - 206 tests passing (91 unit + 115 integration) - Coverage: 58.79% on core modules - auth.py: 57% → 85%, litellm.py: 23% → 87%, chat_summary.py: 41% → 100% - auth_service: 96.51%, config: 100%, rbac: 93.48%
113 lines
4.1 KiB
Python
113 lines
4.1 KiB
Python
"""
|
|
Integration tests for TOTP 2FA operations.
|
|
"""
|
|
|
|
import pyotp
|
|
|
|
|
|
class TestTOTPSetup:
|
|
"""Test TOTP 2FA setup endpoints."""
|
|
|
|
def test_generate_2fa_secret(self, test_app, admin_headers):
|
|
"""Test generating a new 2FA secret."""
|
|
response = test_app.post("/api/auth/setup-2fa/generate", headers=admin_headers)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "secret" in data
|
|
assert "uri" in data
|
|
assert len(data["secret"]) == 32 # Base32 secret length
|
|
assert "otpauth://totp/" in data["uri"]
|
|
assert "NAS" in data["uri"] and "Dashboard" in data["uri"]
|
|
|
|
def test_generate_2fa_without_auth(self, test_app):
|
|
"""Test generating 2FA secret without authentication."""
|
|
response = test_app.post("/api/auth/setup-2fa/generate")
|
|
|
|
assert response.status_code == 401
|
|
|
|
def test_verify_2fa_setup_success(self, test_app, admin_headers):
|
|
"""Test verifying and enabling 2FA with valid code."""
|
|
# First generate a secret
|
|
gen_response = test_app.post("/api/auth/setup-2fa/generate", headers=admin_headers)
|
|
secret = gen_response.json()["secret"]
|
|
|
|
# Generate a valid TOTP code
|
|
totp = pyotp.TOTP(secret)
|
|
valid_code = totp.now()
|
|
|
|
# Verify the code
|
|
response = test_app.post(
|
|
"/api/auth/setup-2fa/verify", headers=admin_headers, json={"secret": secret, "code": valid_code}
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "message" in data
|
|
assert "enabled" in data["message"].lower()
|
|
|
|
def test_verify_2fa_setup_invalid_code(self, test_app, admin_headers):
|
|
"""Test verifying 2FA with invalid code."""
|
|
response = test_app.post(
|
|
"/api/auth/setup-2fa/verify", headers=admin_headers, json={"secret": "JBSWY3DPEHPK3PXP", "code": "000000"}
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
data = response.json()
|
|
assert "detail" in data
|
|
assert "invalid" in data["detail"].lower()
|
|
|
|
def test_verify_2fa_without_auth(self, test_app):
|
|
"""Test verifying 2FA without authentication."""
|
|
response = test_app.post("/api/auth/setup-2fa/verify", json={"secret": "JBSWY3DPEHPK3PXP", "code": "123456"})
|
|
|
|
assert response.status_code == 401
|
|
|
|
def test_disable_2fa(self, test_app, admin_headers):
|
|
"""Test disabling 2FA."""
|
|
response = test_app.post("/api/auth/setup-2fa/disable", headers=admin_headers)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "message" in data
|
|
assert "disabled" in data["message"].lower()
|
|
|
|
def test_disable_2fa_without_auth(self, test_app):
|
|
"""Test disabling 2FA without authentication."""
|
|
response = test_app.post("/api/auth/setup-2fa/disable")
|
|
|
|
assert response.status_code == 401
|
|
|
|
|
|
class TestTOTPWorkflow:
|
|
"""Test complete 2FA setup workflow."""
|
|
|
|
def test_complete_2fa_workflow(self, test_app, admin_headers, mock_config, temp_auth_file):
|
|
"""Test the complete 2FA setup and disable workflow."""
|
|
# Step 1: Generate secret
|
|
gen_response = test_app.post("/api/auth/setup-2fa/generate", headers=admin_headers)
|
|
assert gen_response.status_code == 200
|
|
secret = gen_response.json()["secret"]
|
|
|
|
# Step 2: Verify and enable with valid code
|
|
totp = pyotp.TOTP(secret)
|
|
valid_code = totp.now()
|
|
verify_response = test_app.post(
|
|
"/api/auth/setup-2fa/verify", headers=admin_headers, json={"secret": secret, "code": valid_code}
|
|
)
|
|
assert verify_response.status_code == 200
|
|
|
|
# Step 3: Verify 2FA is enabled by checking auth_service
|
|
from auth_service import load_totp_secret
|
|
|
|
saved_secret = load_totp_secret()
|
|
assert saved_secret == secret
|
|
|
|
# Step 4: Disable 2FA
|
|
disable_response = test_app.post("/api/auth/setup-2fa/disable", headers=admin_headers)
|
|
assert disable_response.status_code == 200
|
|
|
|
# Step 5: Verify 2FA is disabled
|
|
saved_secret = load_totp_secret()
|
|
assert saved_secret == ""
|