chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,272 @@
|
||||
import importlib
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import bcrypt
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from starlette.status import HTTP_403_FORBIDDEN
|
||||
|
||||
from lightrag.api.passwords import BCRYPT_PASSWORD_PREFIX, hash_password
|
||||
from lightrag.tools.hash_password import main as hash_password_main
|
||||
from lightrag.utils import logger as lightrag_logger
|
||||
|
||||
|
||||
def import_real_api_module(module_name: str):
|
||||
sys.modules.pop(module_name, None)
|
||||
|
||||
package_name, _, child_name = module_name.rpartition(".")
|
||||
package = sys.modules.get(package_name)
|
||||
if package is not None and hasattr(package, child_name):
|
||||
delattr(package, child_name)
|
||||
|
||||
return importlib.import_module(module_name)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def auth_module(monkeypatch):
|
||||
config = import_real_api_module("lightrag.api.config")
|
||||
|
||||
mock_global_args = SimpleNamespace(
|
||||
token_secret="test-jwt-secret",
|
||||
jwt_algorithm="HS256",
|
||||
token_expire_hours=48,
|
||||
guest_token_expire_hours=24,
|
||||
auth_accounts="admin:admin_pass",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(config, "global_args", mock_global_args)
|
||||
|
||||
module = import_real_api_module("lightrag.api.auth")
|
||||
module = importlib.reload(module)
|
||||
yield module
|
||||
sys.modules.pop("lightrag.api.auth", None)
|
||||
|
||||
|
||||
def build_bcrypt_value(password: str) -> str:
|
||||
hashed = bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt()).decode("utf-8")
|
||||
return f"{BCRYPT_PASSWORD_PREFIX}{hashed}"
|
||||
|
||||
|
||||
def test_verify_plaintext_password(auth_module):
|
||||
handler = auth_module.AuthHandler()
|
||||
handler.accounts = {"admin": "admin_pass"}
|
||||
|
||||
assert handler.verify_password("admin", "admin_pass")
|
||||
assert not handler.verify_password("admin", "wrong_pass")
|
||||
|
||||
|
||||
def test_verify_prefixed_bcrypt_password(auth_module):
|
||||
handler = auth_module.AuthHandler()
|
||||
handler.accounts = {"user": build_bcrypt_value("user_pass")}
|
||||
|
||||
assert handler.verify_password("user", "user_pass")
|
||||
assert not handler.verify_password("user", "wrong_pass")
|
||||
|
||||
|
||||
def test_plaintext_password_with_bcrypt_prefix_stays_plaintext(auth_module):
|
||||
handler = auth_module.AuthHandler()
|
||||
handler.accounts = {"user": "$2b$not-a-real-hash"}
|
||||
|
||||
assert handler.verify_password("user", "$2b$not-a-real-hash")
|
||||
assert not handler.verify_password("user", "anything-else")
|
||||
|
||||
|
||||
def test_invalid_auth_accounts_raises(monkeypatch):
|
||||
config = import_real_api_module("lightrag.api.config")
|
||||
|
||||
mock_global_args = SimpleNamespace(
|
||||
token_secret="test-jwt-secret",
|
||||
jwt_algorithm="HS256",
|
||||
token_expire_hours=48,
|
||||
guest_token_expire_hours=24,
|
||||
auth_accounts="admin",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(config, "global_args", mock_global_args)
|
||||
|
||||
with pytest.raises(ValueError, match="AUTH_ACCOUNTS must use"):
|
||||
import_real_api_module("lightrag.api.auth")
|
||||
|
||||
sys.modules.pop("lightrag.api.auth", None)
|
||||
|
||||
|
||||
def test_initialize_config_rejects_default_token_secret_with_auth_accounts():
|
||||
config = import_real_api_module("lightrag.api.config")
|
||||
|
||||
insecure_args = SimpleNamespace(
|
||||
auth_accounts="admin:admin_pass",
|
||||
token_secret=config.DEFAULT_TOKEN_SECRET,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="TOKEN_SECRET must be explicitly set"):
|
||||
config.initialize_config(insecure_args, force=True)
|
||||
|
||||
|
||||
def test_initialize_config_allows_custom_token_secret_with_auth_accounts():
|
||||
config = import_real_api_module("lightrag.api.config")
|
||||
|
||||
secure_args = SimpleNamespace(
|
||||
auth_accounts="admin:admin_pass",
|
||||
token_secret="custom-jwt-secret",
|
||||
)
|
||||
|
||||
initialized = config.initialize_config(secure_args, force=True)
|
||||
|
||||
assert initialized is secure_args
|
||||
|
||||
|
||||
def test_guest_tokens_fall_back_to_default_secret_when_token_secret_missing(
|
||||
monkeypatch,
|
||||
):
|
||||
config = import_real_api_module("lightrag.api.config")
|
||||
|
||||
mock_global_args = SimpleNamespace(
|
||||
token_secret=None,
|
||||
jwt_algorithm="HS256",
|
||||
token_expire_hours=48,
|
||||
guest_token_expire_hours=24,
|
||||
auth_accounts="",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(config, "global_args", mock_global_args)
|
||||
warning_messages = []
|
||||
|
||||
def capture_warning(message):
|
||||
warning_messages.append(message)
|
||||
|
||||
monkeypatch.setattr(lightrag_logger, "warning", capture_warning)
|
||||
|
||||
module = import_real_api_module("lightrag.api.auth")
|
||||
module = importlib.reload(module)
|
||||
handler = module.AuthHandler()
|
||||
|
||||
token = handler.create_token("guest", role="guest")
|
||||
token_info = handler.validate_token(token)
|
||||
|
||||
assert handler.secret == config.DEFAULT_TOKEN_SECRET
|
||||
assert token_info["username"] == "guest"
|
||||
assert token_info["role"] == "guest"
|
||||
assert any(
|
||||
"Falling back to the default guest-mode JWT secret" in msg
|
||||
for msg in warning_messages
|
||||
)
|
||||
|
||||
sys.modules.pop("lightrag.api.auth", None)
|
||||
|
||||
|
||||
def _load_api_key_only_modules(monkeypatch):
|
||||
"""Import auth + utils_api in the API-key-only profile.
|
||||
|
||||
API-key-only = an API key is set but AUTH_ACCOUNTS is empty, so
|
||||
``auth_configured`` is False and AuthHandler falls back to the public
|
||||
DEFAULT_TOKEN_SECRET. Returns (config, auth, utils_api).
|
||||
"""
|
||||
config = import_real_api_module("lightrag.api.config")
|
||||
|
||||
mock_global_args = SimpleNamespace(
|
||||
token_secret=None, # -> AuthHandler falls back to DEFAULT_TOKEN_SECRET
|
||||
jwt_algorithm="HS256",
|
||||
token_expire_hours=48,
|
||||
guest_token_expire_hours=24,
|
||||
auth_accounts="", # no password accounts -> auth_configured is False
|
||||
whitelist_paths="/health", # consumed by utils_api at import time
|
||||
token_auto_renew=False,
|
||||
)
|
||||
monkeypatch.setattr(config, "global_args", mock_global_args)
|
||||
|
||||
auth = import_real_api_module("lightrag.api.auth")
|
||||
auth = importlib.reload(auth)
|
||||
utils_api = import_real_api_module("lightrag.api.utils_api")
|
||||
return config, auth, utils_api
|
||||
|
||||
|
||||
async def test_combined_auth_rejects_guest_token_in_api_key_mode(monkeypatch):
|
||||
"""A forged/obtained guest token must not bypass the X-API-Key check.
|
||||
|
||||
Regression for GHSA-f4vv-55c2-5789 / GHSA-xr5c-v5r6-c9f9: in the
|
||||
API-key-only profile an anonymous caller can mint a guest JWT with the
|
||||
public default secret (or fetch one from /auth-status) and previously
|
||||
short-circuited authorization before the API key was ever checked.
|
||||
"""
|
||||
config, auth, utils_api = _load_api_key_only_modules(monkeypatch)
|
||||
try:
|
||||
assert utils_api.auth_configured is False
|
||||
assert auth.auth_handler.secret == config.DEFAULT_TOKEN_SECRET
|
||||
|
||||
forged_guest = auth.auth_handler.create_token(username="guest", role="guest")
|
||||
|
||||
api_key = "operator-secret-key"
|
||||
dependency = utils_api.get_combined_auth_dependency(api_key=api_key)
|
||||
|
||||
request = SimpleNamespace(url=SimpleNamespace(path="/documents"), scope={})
|
||||
response = SimpleNamespace(headers={})
|
||||
|
||||
# Forged guest token, no API key -> rejected (was HTTP 200 before the fix).
|
||||
with pytest.raises(HTTPException) as rejected:
|
||||
await dependency(
|
||||
request=request,
|
||||
response=response,
|
||||
token=forged_guest,
|
||||
api_key_header_value=None,
|
||||
)
|
||||
assert rejected.value.status_code == HTTP_403_FORBIDDEN
|
||||
|
||||
# Guest token together with the real API key -> allowed (the WebUI sends
|
||||
# both Authorization and X-API-Key headers).
|
||||
result = await dependency(
|
||||
request=request,
|
||||
response=response,
|
||||
token=forged_guest,
|
||||
api_key_header_value=api_key,
|
||||
)
|
||||
assert result is None
|
||||
finally:
|
||||
sys.modules.pop("lightrag.api.utils_api", None)
|
||||
sys.modules.pop("lightrag.api.auth", None)
|
||||
|
||||
|
||||
async def test_combined_auth_allows_guest_token_when_fully_open(monkeypatch):
|
||||
"""Fully-open mode (no API key, no AUTH_ACCOUNTS) keeps accepting guest tokens.
|
||||
|
||||
Guards against over-fixing: the bypass fix must not break zero-config guest
|
||||
access, which is the intended convenience mode.
|
||||
"""
|
||||
_config, auth, utils_api = _load_api_key_only_modules(monkeypatch)
|
||||
try:
|
||||
guest_token = auth.auth_handler.create_token(username="guest", role="guest")
|
||||
|
||||
dependency = utils_api.get_combined_auth_dependency(api_key=None)
|
||||
request = SimpleNamespace(url=SimpleNamespace(path="/documents"), scope={})
|
||||
response = SimpleNamespace(headers={})
|
||||
|
||||
result = await dependency(
|
||||
request=request,
|
||||
response=response,
|
||||
token=guest_token,
|
||||
api_key_header_value=None,
|
||||
)
|
||||
assert result is None
|
||||
finally:
|
||||
sys.modules.pop("lightrag.api.utils_api", None)
|
||||
sys.modules.pop("lightrag.api.auth", None)
|
||||
|
||||
|
||||
def test_hash_password_returns_prefixed_value(auth_module):
|
||||
hashed = hash_password("new_password")
|
||||
|
||||
assert hashed.startswith(BCRYPT_PASSWORD_PREFIX)
|
||||
raw_hash = hashed[len(BCRYPT_PASSWORD_PREFIX) :]
|
||||
assert bcrypt.checkpw("new_password".encode("utf-8"), raw_hash.encode("utf-8"))
|
||||
|
||||
|
||||
def test_hash_password_cli_outputs_auth_accounts_entry(capsys):
|
||||
exit_code = hash_password_main(["--username", "admin", "secret"])
|
||||
|
||||
assert exit_code == 0
|
||||
output = capsys.readouterr().out.strip()
|
||||
username, hashed = output.split(":", 1)
|
||||
assert username == "admin"
|
||||
assert hashed.startswith(BCRYPT_PASSWORD_PREFIX)
|
||||
raw_hash = hashed[len(BCRYPT_PASSWORD_PREFIX) :]
|
||||
assert bcrypt.checkpw("secret".encode("utf-8"), raw_hash.encode("utf-8"))
|
||||
@@ -0,0 +1,403 @@
|
||||
"""
|
||||
Pytest unit tests for token auto-renewal functionality
|
||||
|
||||
Tests:
|
||||
1. Backend token renewal logic
|
||||
2. Rate limiting for token renewals
|
||||
3. Token renewal state tracking
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import Mock
|
||||
from fastapi import Response
|
||||
import time
|
||||
|
||||
# Create a simple token renewal cache for testing
|
||||
_token_renewal_cache = {}
|
||||
_RENEWAL_MIN_INTERVAL = 60
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
class TestTokenRenewal:
|
||||
"""Tests for token auto-renewal logic"""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_auth_handler(self):
|
||||
"""Mock authentication handler"""
|
||||
handler = Mock()
|
||||
handler.guest_expire_hours = 24
|
||||
handler.expire_hours = 24
|
||||
handler.create_token = Mock(return_value="new-token-12345")
|
||||
return handler
|
||||
|
||||
@pytest.fixture
|
||||
def mock_global_args(self):
|
||||
"""Mock global configuration"""
|
||||
args = Mock()
|
||||
args.token_auto_renew = True
|
||||
args.token_renew_threshold = 0.5
|
||||
return args
|
||||
|
||||
@pytest.fixture
|
||||
def mock_token_info_guest(self):
|
||||
"""Mock token info for guest user"""
|
||||
# Token with 10 hours remaining (below 50% of 24 hours)
|
||||
exp_time = datetime.now(timezone.utc) + timedelta(hours=10)
|
||||
return {
|
||||
"username": "guest",
|
||||
"role": "guest",
|
||||
"exp": exp_time,
|
||||
"metadata": {"auth_mode": "disabled"},
|
||||
}
|
||||
|
||||
@pytest.fixture
|
||||
def mock_token_info_user(self):
|
||||
"""Mock token info for regular user"""
|
||||
# Token with 10 hours remaining (below 50% of 24 hours)
|
||||
exp_time = datetime.now(timezone.utc) + timedelta(hours=10)
|
||||
return {
|
||||
"username": "testuser",
|
||||
"role": "user",
|
||||
"exp": exp_time,
|
||||
"metadata": {"auth_mode": "enabled"},
|
||||
}
|
||||
|
||||
@pytest.fixture
|
||||
def mock_token_info_above_threshold(self):
|
||||
"""Mock token info with time above renewal threshold"""
|
||||
# Token with 20 hours remaining (above 50% of 24 hours)
|
||||
exp_time = datetime.now(timezone.utc) + timedelta(hours=20)
|
||||
return {
|
||||
"username": "testuser",
|
||||
"role": "user",
|
||||
"exp": exp_time,
|
||||
"metadata": {"auth_mode": "enabled"},
|
||||
}
|
||||
|
||||
def test_token_renewal_when_below_threshold(
|
||||
self, mock_auth_handler, mock_global_args, mock_token_info_user
|
||||
):
|
||||
"""Test that token is renewed when remaining time < threshold"""
|
||||
# Use global cache
|
||||
global _token_renewal_cache
|
||||
|
||||
# Clear cache
|
||||
_token_renewal_cache.clear()
|
||||
|
||||
response = Mock(spec=Response)
|
||||
response.headers = {}
|
||||
|
||||
# Simulate the renewal logic
|
||||
expire_time = mock_token_info_user["exp"]
|
||||
now = datetime.now(timezone.utc)
|
||||
remaining_seconds = (expire_time - now).total_seconds()
|
||||
|
||||
role = mock_token_info_user["role"]
|
||||
total_hours = (
|
||||
mock_auth_handler.expire_hours
|
||||
if role == "user"
|
||||
else mock_auth_handler.guest_expire_hours
|
||||
)
|
||||
total_seconds = total_hours * 3600
|
||||
|
||||
# Should renew because remaining_seconds < total_seconds * 0.5
|
||||
should_renew = (
|
||||
remaining_seconds < total_seconds * mock_global_args.token_renew_threshold
|
||||
)
|
||||
assert should_renew is True
|
||||
|
||||
# Simulate renewal
|
||||
username = mock_token_info_user["username"]
|
||||
current_time = time.time()
|
||||
last_renewal = _token_renewal_cache.get(username, 0)
|
||||
time_since_last_renewal = current_time - last_renewal
|
||||
|
||||
# Should pass rate limit (first renewal)
|
||||
assert time_since_last_renewal >= 60 or last_renewal == 0
|
||||
|
||||
# Perform renewal
|
||||
new_token = mock_auth_handler.create_token(
|
||||
username=username, role=role, metadata=mock_token_info_user["metadata"]
|
||||
)
|
||||
response.headers["X-New-Token"] = new_token
|
||||
_token_renewal_cache[username] = current_time
|
||||
|
||||
# Verify
|
||||
assert "X-New-Token" in response.headers
|
||||
assert response.headers["X-New-Token"] == "new-token-12345"
|
||||
assert username in _token_renewal_cache
|
||||
|
||||
def test_token_no_renewal_when_above_threshold(
|
||||
self, mock_auth_handler, mock_global_args, mock_token_info_above_threshold
|
||||
):
|
||||
"""Test that token is NOT renewed when remaining time > threshold"""
|
||||
response = Mock(spec=Response)
|
||||
response.headers = {}
|
||||
|
||||
expire_time = mock_token_info_above_threshold["exp"]
|
||||
now = datetime.now(timezone.utc)
|
||||
remaining_seconds = (expire_time - now).total_seconds()
|
||||
|
||||
mock_token_info_above_threshold["role"]
|
||||
total_hours = mock_auth_handler.expire_hours
|
||||
total_seconds = total_hours * 3600
|
||||
|
||||
# Should NOT renew because remaining_seconds > total_seconds * 0.5
|
||||
should_renew = (
|
||||
remaining_seconds < total_seconds * mock_global_args.token_renew_threshold
|
||||
)
|
||||
assert should_renew is False
|
||||
|
||||
# No renewal should happen
|
||||
assert "X-New-Token" not in response.headers
|
||||
|
||||
def test_token_renewal_disabled(
|
||||
self, mock_auth_handler, mock_global_args, mock_token_info_user
|
||||
):
|
||||
"""Test that no renewal happens when TOKEN_AUTO_RENEW=false"""
|
||||
mock_global_args.token_auto_renew = False
|
||||
response = Mock(spec=Response)
|
||||
response.headers = {}
|
||||
|
||||
# Auto-renewal is disabled, so even if below threshold, no renewal
|
||||
if not mock_global_args.token_auto_renew:
|
||||
# Skip renewal logic
|
||||
pass
|
||||
|
||||
assert "X-New-Token" not in response.headers
|
||||
|
||||
def test_token_renewal_for_guest_mode(
|
||||
self, mock_auth_handler, mock_global_args, mock_token_info_guest
|
||||
):
|
||||
"""Test that guest tokens are renewed correctly"""
|
||||
# Use global cache
|
||||
global _token_renewal_cache
|
||||
|
||||
_token_renewal_cache.clear()
|
||||
|
||||
response = Mock(spec=Response)
|
||||
response.headers = {}
|
||||
|
||||
expire_time = mock_token_info_guest["exp"]
|
||||
now = datetime.now(timezone.utc)
|
||||
remaining_seconds = (expire_time - now).total_seconds()
|
||||
|
||||
role = mock_token_info_guest["role"]
|
||||
total_hours = mock_auth_handler.guest_expire_hours
|
||||
total_seconds = total_hours * 3600
|
||||
|
||||
should_renew = (
|
||||
remaining_seconds < total_seconds * mock_global_args.token_renew_threshold
|
||||
)
|
||||
assert should_renew is True
|
||||
|
||||
# Renewal for guest
|
||||
username = mock_token_info_guest["username"]
|
||||
new_token = mock_auth_handler.create_token(
|
||||
username=username, role=role, metadata=mock_token_info_guest["metadata"]
|
||||
)
|
||||
response.headers["X-New-Token"] = new_token
|
||||
_token_renewal_cache[username] = time.time()
|
||||
|
||||
assert "X-New-Token" in response.headers
|
||||
assert username in _token_renewal_cache
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
class TestRateLimiting:
|
||||
"""Tests for token renewal rate limiting"""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_auth_handler(self):
|
||||
"""Mock authentication handler"""
|
||||
handler = Mock()
|
||||
handler.expire_hours = 24
|
||||
handler.create_token = Mock(return_value="new-token-12345")
|
||||
return handler
|
||||
|
||||
def test_rate_limit_prevents_rapid_renewals(self, mock_auth_handler):
|
||||
"""Test that second renewal within 60s is blocked"""
|
||||
# Use global cache and constant
|
||||
global _token_renewal_cache, _RENEWAL_MIN_INTERVAL
|
||||
|
||||
username = "testuser"
|
||||
_token_renewal_cache.clear()
|
||||
|
||||
# First renewal
|
||||
current_time_1 = time.time()
|
||||
_token_renewal_cache[username] = current_time_1
|
||||
|
||||
response_1 = Mock(spec=Response)
|
||||
response_1.headers = {}
|
||||
response_1.headers["X-New-Token"] = "new-token-12345"
|
||||
|
||||
# Immediate second renewal attempt (within 60s)
|
||||
current_time_2 = time.time() # Almost same time
|
||||
last_renewal = _token_renewal_cache.get(username, 0)
|
||||
time_since_last_renewal = current_time_2 - last_renewal
|
||||
|
||||
# Should be blocked by rate limit
|
||||
assert time_since_last_renewal < _RENEWAL_MIN_INTERVAL
|
||||
|
||||
response_2 = Mock(spec=Response)
|
||||
response_2.headers = {}
|
||||
|
||||
# No new token should be issued
|
||||
if time_since_last_renewal < _RENEWAL_MIN_INTERVAL:
|
||||
# Rate limited, skip renewal
|
||||
pass
|
||||
|
||||
assert "X-New-Token" not in response_2.headers
|
||||
|
||||
def test_rate_limit_allows_renewal_after_interval(self, mock_auth_handler):
|
||||
"""Test that renewal succeeds after 60s interval"""
|
||||
# Use global cache and constant
|
||||
global _token_renewal_cache, _RENEWAL_MIN_INTERVAL
|
||||
|
||||
username = "testuser"
|
||||
_token_renewal_cache.clear()
|
||||
|
||||
# First renewal at time T
|
||||
first_renewal_time = time.time() - 61 # 61 seconds ago
|
||||
_token_renewal_cache[username] = first_renewal_time
|
||||
|
||||
# Second renewal attempt now
|
||||
current_time = time.time()
|
||||
last_renewal = _token_renewal_cache.get(username, 0)
|
||||
time_since_last_renewal = current_time - last_renewal
|
||||
|
||||
# Should pass rate limit (>60s elapsed)
|
||||
assert time_since_last_renewal >= _RENEWAL_MIN_INTERVAL
|
||||
|
||||
response = Mock(spec=Response)
|
||||
response.headers = {}
|
||||
|
||||
if time_since_last_renewal >= _RENEWAL_MIN_INTERVAL:
|
||||
new_token = mock_auth_handler.create_token(
|
||||
username=username, role="user", metadata={}
|
||||
)
|
||||
response.headers["X-New-Token"] = new_token
|
||||
_token_renewal_cache[username] = current_time
|
||||
|
||||
assert "X-New-Token" in response.headers
|
||||
assert response.headers["X-New-Token"] == "new-token-12345"
|
||||
|
||||
def test_rate_limit_per_user(self, mock_auth_handler):
|
||||
"""Test that different users have independent rate limits"""
|
||||
# Use global cache
|
||||
global _token_renewal_cache
|
||||
|
||||
_token_renewal_cache.clear()
|
||||
|
||||
user1 = "user1"
|
||||
user2 = "user2"
|
||||
|
||||
current_time = time.time()
|
||||
|
||||
# User1 gets renewal
|
||||
_token_renewal_cache[user1] = current_time
|
||||
|
||||
# User2 should still be able to get renewal (independent cache)
|
||||
last_renewal_user2 = _token_renewal_cache.get(user2, 0)
|
||||
assert last_renewal_user2 == 0 # No previous renewal
|
||||
|
||||
# User2 can renew
|
||||
_token_renewal_cache[user2] = current_time
|
||||
|
||||
# Both users should have entries
|
||||
assert user1 in _token_renewal_cache
|
||||
assert user2 in _token_renewal_cache
|
||||
assert _token_renewal_cache[user1] == _token_renewal_cache[user2]
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
class TestTokenExpirationCalculation:
|
||||
"""Tests for token expiration time calculation"""
|
||||
|
||||
def test_expiration_extraction_from_jwt(self):
|
||||
"""Test extracting expiration time from JWT token"""
|
||||
import base64
|
||||
import json
|
||||
|
||||
# Create a mock JWT payload
|
||||
exp_timestamp = int(
|
||||
(datetime.now(timezone.utc) + timedelta(hours=24)).timestamp()
|
||||
)
|
||||
payload = {"sub": "testuser", "role": "user", "exp": exp_timestamp}
|
||||
|
||||
# Encode as base64 (simulating JWT structure: header.payload.signature)
|
||||
payload_b64 = base64.b64encode(json.dumps(payload).encode()).decode()
|
||||
mock_token = f"header.{payload_b64}.signature"
|
||||
|
||||
# Simulate extraction
|
||||
parts = mock_token.split(".")
|
||||
assert len(parts) == 3
|
||||
|
||||
decoded_payload = json.loads(base64.b64decode(parts[1]))
|
||||
assert decoded_payload["exp"] == exp_timestamp
|
||||
assert decoded_payload["sub"] == "testuser"
|
||||
|
||||
def test_remaining_time_calculation(self):
|
||||
"""Test calculation of remaining token time"""
|
||||
# Token expires in 10 hours
|
||||
exp_time = datetime.now(timezone.utc) + timedelta(hours=10)
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
remaining_seconds = (exp_time - now).total_seconds()
|
||||
|
||||
# Should be approximately 10 hours (36000 seconds)
|
||||
assert 35990 < remaining_seconds < 36010
|
||||
|
||||
# Calculate percentage remaining (for 24-hour token)
|
||||
total_seconds = 24 * 3600
|
||||
percentage_remaining = remaining_seconds / total_seconds
|
||||
|
||||
# Should be approximately 41.67% remaining
|
||||
assert 0.41 < percentage_remaining < 0.42
|
||||
|
||||
def test_threshold_comparison(self):
|
||||
"""Test threshold-based renewal decision"""
|
||||
threshold = 0.5
|
||||
total_hours = 24
|
||||
total_seconds = total_hours * 3600
|
||||
|
||||
# Scenario 1: 10 hours remaining -> should renew
|
||||
remaining_seconds_1 = 10 * 3600
|
||||
should_renew_1 = remaining_seconds_1 < total_seconds * threshold
|
||||
assert should_renew_1 is True
|
||||
|
||||
# Scenario 2: 20 hours remaining -> should NOT renew
|
||||
remaining_seconds_2 = 20 * 3600
|
||||
should_renew_2 = remaining_seconds_2 < total_seconds * threshold
|
||||
assert should_renew_2 is False
|
||||
|
||||
# Scenario 3: Exactly 12 hours remaining (at threshold) -> should NOT renew
|
||||
remaining_seconds_3 = 12 * 3600
|
||||
should_renew_3 = remaining_seconds_3 < total_seconds * threshold
|
||||
assert should_renew_3 is False
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
def test_renewal_cache_cleanup():
|
||||
"""Test that renewal cache can be cleared"""
|
||||
# Use global cache
|
||||
global _token_renewal_cache
|
||||
|
||||
# Clear first
|
||||
_token_renewal_cache.clear()
|
||||
|
||||
# Add some entries
|
||||
_token_renewal_cache["user1"] = time.time()
|
||||
_token_renewal_cache["user2"] = time.time()
|
||||
|
||||
assert len(_token_renewal_cache) == 2
|
||||
|
||||
# Clear cache
|
||||
_token_renewal_cache.clear()
|
||||
|
||||
assert len(_token_renewal_cache) == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
Reference in New Issue
Block a user