250 lines
10 KiB
Python
250 lines
10 KiB
Python
import pytest
|
|
import jwt
|
|
import time
|
|
import uuid
|
|
from unittest.mock import MagicMock, patch
|
|
from datetime import datetime, UTC
|
|
from fastapi import HTTPException
|
|
|
|
# Mock environment variable before importing auth
|
|
with patch.dict('os.environ', {'JWT_AAREDAQ_KEY': 'test_secret'}):
|
|
from aare.daq.auth import (
|
|
TokenData, create_access_token, authenticate_user, parse_token,
|
|
check_jwt_ro, check_jwt_rw, check_jwt_staff_only, check_jwt_staff,
|
|
force_current_sesion, get_baton_status, request_baton,
|
|
respond_to_baton_request, release_baton, cancel_baton_request,
|
|
resolve_baton_timeout_if_needed
|
|
)
|
|
from aare.common.exception_handler import AuthenticationException, UserRightsException
|
|
from aare.common.auth_models import BatonStatus, BatonRequest, BatonRequestStatus, BatonHolderInfo, BatonTransferQueue
|
|
from aare.common.models import SessionsStateEnum
|
|
|
|
@pytest.fixture
|
|
def mock_cfg():
|
|
cfg = MagicMock()
|
|
cfg.pgroup = "p12345"
|
|
cfg.generate_session.return_value = 100
|
|
cfg.queued_baton_transfer = None
|
|
cfg.baton_holder = None
|
|
cfg.pending_baton_request = None
|
|
return cfg
|
|
|
|
@pytest.fixture
|
|
def token_data():
|
|
return TokenData(sub="testuser", pgroups=["p12345", "p67890"], session=100, staff=False)
|
|
|
|
@pytest.fixture
|
|
def staff_token_data():
|
|
return TokenData(sub="staffuser", pgroups=["p12345"], session=101, staff=True)
|
|
|
|
def test_create_access_token(token_data):
|
|
with patch('aare.daq.auth.SECRET_KEY', 'test_secret'):
|
|
token = create_access_token(token_data)
|
|
assert isinstance(token, str)
|
|
payload = jwt.decode(token, 'test_secret', algorithms=["HS256"])
|
|
assert payload["sub"] == "testuser"
|
|
assert payload["session"] == 100
|
|
|
|
def test_authenticate_user(mock_cfg):
|
|
with patch('pwd.getpwnam') as mock_pwd, \
|
|
patch('os.getgrouplist') as mock_groups, \
|
|
patch('grp.getgrgid') as mock_grp, \
|
|
patch('aare.daq.auth.SECRET_KEY', 'test_secret'):
|
|
|
|
mock_pwd.return_value.pw_name = "testuser"
|
|
mock_pwd.return_value.pw_gid = 1000
|
|
mock_groups.return_value = [1000, 1001]
|
|
|
|
def get_group(gid):
|
|
m = MagicMock()
|
|
if gid == 1000: m.gr_name = "p12345"
|
|
else: m.gr_name = "unx-MXgroup"
|
|
return m
|
|
|
|
mock_grp.side_effect = get_group
|
|
|
|
form_data = MagicMock()
|
|
form_data.username = "testuser"
|
|
|
|
token = authenticate_user(mock_cfg, form_data)
|
|
assert isinstance(token, str)
|
|
|
|
payload = jwt.decode(token, 'test_secret', algorithms=["HS256"])
|
|
assert payload["sub"] == "testuser"
|
|
assert "p12345" in payload["pgroups"]
|
|
assert payload["staff"] is True
|
|
|
|
def test_parse_token():
|
|
with patch('aare.daq.auth.SECRET_KEY', 'test_secret'):
|
|
token = jwt.encode({"sub": "user", "pgroups": [], "session": 1, "staff": False}, "test_secret")
|
|
data = parse_token(token)
|
|
assert data.sub == "user"
|
|
|
|
def test_parse_token_invalid():
|
|
with pytest.raises(AuthenticationException):
|
|
parse_token("invalid.token.here")
|
|
|
|
def test_check_jwt_ro(mock_cfg, token_data):
|
|
# Success
|
|
check_jwt_ro(mock_cfg, token_data)
|
|
|
|
# Fail
|
|
mock_cfg.pgroup = "p99999"
|
|
with pytest.raises(UserRightsException):
|
|
check_jwt_ro(mock_cfg, token_data)
|
|
|
|
# Staff success even if not in pgroup
|
|
staff_data = TokenData(sub="staff", pgroups=[], session=2, staff=True)
|
|
check_jwt_ro(mock_cfg, staff_data)
|
|
|
|
def test_check_jwt_rw(mock_cfg, token_data):
|
|
mock_cfg.baton_holder = BatonHolderInfo(session=100, username="testuser", is_staff=False)
|
|
check_jwt_rw(mock_cfg, token_data)
|
|
|
|
# Not holder
|
|
mock_cfg.baton_holder = BatonHolderInfo(session=200, username="other", is_staff=False)
|
|
with pytest.raises(UserRightsException):
|
|
check_jwt_rw(mock_cfg, token_data)
|
|
|
|
def test_check_jwt_staff_only(token_data, staff_token_data):
|
|
check_jwt_staff_only(staff_token_data)
|
|
with pytest.raises(UserRightsException):
|
|
check_jwt_staff_only(token_data)
|
|
|
|
def test_force_current_sesion(mock_cfg, token_data):
|
|
force_current_sesion(mock_cfg, token_data)
|
|
mock_cfg.execute_baton_transfer.assert_called_once()
|
|
|
|
def test_get_baton_status(mock_cfg, token_data):
|
|
mock_cfg.baton_holder = None
|
|
mock_cfg.pending_baton_request = None
|
|
mock_cfg.queued_baton_transfer = None
|
|
|
|
status = get_baton_status(mock_cfg, token_data)
|
|
assert isinstance(status, BatonStatus)
|
|
assert status.you_are_holder is False
|
|
|
|
def test_request_baton_vacant(mock_cfg, token_data):
|
|
mock_cfg.session_state.return_value = SessionsStateEnum.Vacant
|
|
res = request_baton(mock_cfg, token_data)
|
|
assert res["granted"] is True
|
|
mock_cfg.execute_baton_transfer.assert_called_once()
|
|
|
|
def test_request_baton_owned_by_you(mock_cfg, token_data):
|
|
mock_cfg.session_state.return_value = SessionsStateEnum.OwnedByYou
|
|
res = request_baton(mock_cfg, token_data)
|
|
assert res["already_holder"] is True
|
|
|
|
def test_request_baton_pending(mock_cfg, token_data):
|
|
mock_cfg.session_state.return_value = SessionsStateEnum.OwnedByElse
|
|
mock_cfg.baton_holder = BatonHolderInfo(session=200, username="other", is_staff=False)
|
|
mock_cfg.pending_baton_request = None
|
|
mock_cfg.can_transfer_baton_now.return_value = True
|
|
mock_cfg.allow_non_staff_request_from_staff = True
|
|
|
|
res = request_baton(mock_cfg, token_data)
|
|
assert res["pending"] is True
|
|
mock_cfg.set_pending_baton_request.assert_called_once()
|
|
|
|
def test_respond_to_baton_request_accept(mock_cfg, token_data):
|
|
mock_cfg.baton_holder = BatonHolderInfo(session=100, username="testuser", is_staff=False)
|
|
mock_cfg.pending_baton_request = BatonRequest(
|
|
request_id="1", requester_username="other", requester_session=200,
|
|
requester_is_staff=False, holder_username="testuser", holder_session=100,
|
|
created_at=time.time(), timeout_seconds=30, status=BatonRequestStatus.PENDING
|
|
)
|
|
mock_cfg.can_transfer_baton_now.return_value = True
|
|
|
|
res = respond_to_baton_request(mock_cfg, token_data, accept=True)
|
|
assert res["transferred"] is True
|
|
mock_cfg.execute_baton_transfer.assert_called_once()
|
|
|
|
def test_release_baton(mock_cfg, token_data):
|
|
mock_cfg.baton_holder = BatonHolderInfo(session=100, username="testuser", is_staff=False)
|
|
res = release_baton(mock_cfg, token_data)
|
|
assert res["released"] is True
|
|
mock_cfg.end_active_session.assert_called_with(100)
|
|
|
|
def test_request_baton_staff_override(mock_cfg, staff_token_data):
|
|
mock_cfg.session_state.return_value = SessionsStateEnum.OwnedByElse
|
|
mock_cfg.baton_holder = BatonHolderInfo(session=200, username="other", is_staff=False)
|
|
mock_cfg.can_transfer_baton_now.return_value = True
|
|
|
|
res = request_baton(mock_cfg, staff_token_data)
|
|
assert res["granted"] is True
|
|
assert res["override"] is True
|
|
|
|
def test_request_baton_staff_override_busy(mock_cfg, staff_token_data):
|
|
mock_cfg.session_state.return_value = SessionsStateEnum.OwnedByElse
|
|
mock_cfg.baton_holder = BatonHolderInfo(session=200, username="other", is_staff=False)
|
|
mock_cfg.can_transfer_baton_now.return_value = False
|
|
|
|
res = request_baton(mock_cfg, staff_token_data)
|
|
assert res["queued"] is True
|
|
|
|
def test_resolve_baton_timeout_if_needed(mock_cfg):
|
|
# Case: No pending request
|
|
mock_cfg.pending_baton_request = None
|
|
assert resolve_baton_timeout_if_needed(mock_cfg) is None
|
|
|
|
# Case: Request not expired
|
|
mock_cfg.pending_baton_request = BatonRequest(
|
|
request_id="1", requester_username="other", requester_session=200,
|
|
requester_is_staff=False, holder_username="testuser", holder_session=100,
|
|
created_at=time.time(), timeout_seconds=30, status=BatonRequestStatus.PENDING
|
|
)
|
|
assert resolve_baton_timeout_if_needed(mock_cfg) is None
|
|
|
|
# Case: Request expired, beamline not busy
|
|
mock_cfg.pending_baton_request.created_at = time.time() - 40
|
|
mock_cfg.can_transfer_baton_now.return_value = True
|
|
mock_cfg.baton_holder = BatonHolderInfo(session=100, username="testuser", is_staff=False)
|
|
|
|
resolve_baton_timeout_if_needed(mock_cfg)
|
|
mock_cfg.execute_baton_transfer.assert_called_once()
|
|
mock_cfg.clear_pending_baton_request.assert_called_once()
|
|
|
|
def test_respond_to_baton_request_refuse(mock_cfg, token_data):
|
|
mock_cfg.baton_holder = BatonHolderInfo(session=100, username="testuser", is_staff=False)
|
|
mock_cfg.pending_baton_request = BatonRequest(
|
|
request_id="1", requester_username="other", requester_session=200,
|
|
requester_is_staff=False, holder_username="testuser", holder_session=100,
|
|
created_at=time.time(), timeout_seconds=30, status=BatonRequestStatus.PENDING
|
|
)
|
|
res = respond_to_baton_request(mock_cfg, token_data, accept=False)
|
|
assert res["refused"] is True
|
|
assert mock_cfg.pending_baton_request.status == BatonRequestStatus.REFUSED
|
|
|
|
def test_respond_to_baton_request_accept_busy(mock_cfg, token_data):
|
|
mock_cfg.baton_holder = BatonHolderInfo(session=100, username="testuser", is_staff=False)
|
|
mock_cfg.pending_baton_request = BatonRequest(
|
|
request_id="1", requester_username="other", requester_session=200,
|
|
requester_is_staff=False, holder_username="testuser", holder_session=100,
|
|
created_at=time.time(), timeout_seconds=30, status=BatonRequestStatus.PENDING
|
|
)
|
|
mock_cfg.can_transfer_baton_now.return_value = False
|
|
|
|
res = respond_to_baton_request(mock_cfg, token_data, accept=True)
|
|
assert res["queued"] is True
|
|
assert mock_cfg.queued_baton_transfer is not None
|
|
|
|
def test_cancel_baton_request(mock_cfg, token_data):
|
|
mock_cfg.pending_baton_request = BatonRequest(
|
|
request_id="1", requester_username="testuser", requester_session=100,
|
|
requester_is_staff=False, holder_username="other", holder_session=200,
|
|
created_at=time.time(), timeout_seconds=30, status=BatonRequestStatus.PENDING
|
|
)
|
|
res = cancel_baton_request(mock_cfg, token_data)
|
|
assert res["cancelled"] is True
|
|
mock_cfg.clear_pending_baton_request.assert_called_once()
|
|
|
|
def test_resolve_baton_timeout_busy(mock_cfg):
|
|
mock_cfg.pending_baton_request = BatonRequest(
|
|
request_id="1", requester_username="other", requester_session=200,
|
|
requester_is_staff=False, holder_username="testuser", holder_session=100,
|
|
created_at=time.time() - 40, timeout_seconds=30, status=BatonRequestStatus.PENDING
|
|
)
|
|
mock_cfg.can_transfer_baton_now.return_value = False
|
|
resolve_baton_timeout_if_needed(mock_cfg)
|
|
assert mock_cfg.queued_baton_transfer is not None
|