iptv-manager-service/tests/auth/test_cognito.py

from unittest.mock import MagicMock, patch

import pytest
from fastapi import HTTPException, status

# Test constants
TEST_CLIENT_ID = "test_client_id"
TEST_CLIENT_SECRET = "test_client_secret"

# Patch constants before importing the module
with (
    patch("app.utils.constants.COGNITO_CLIENT_ID", TEST_CLIENT_ID),
    patch("app.utils.constants.COGNITO_CLIENT_SECRET", TEST_CLIENT_SECRET),
):
    from app.auth.cognito import get_user_from_token, initiate_auth
    from app.models.auth import CognitoUser
    from app.utils.constants import USER_ROLE_ATTRIBUTE


@pytest.fixture(autouse=True)
def mock_cognito_client():
    with patch("app.auth.cognito.cognito_client") as mock_client:
        # Setup mock client and exceptions
        mock_client.exceptions = MagicMock()
        mock_client.exceptions.NotAuthorizedException = type(
            "NotAuthorizedException", (Exception,), {}
        )
        mock_client.exceptions.UserNotFoundException = type(
            "UserNotFoundException", (Exception,), {}
        )
        yield mock_client


def test_initiate_auth_success(mock_cognito_client):
    # Mock successful authentication response
    mock_cognito_client.initiate_auth.return_value = {
        "AuthenticationResult": {
            "AccessToken": "mock_access_token",
            "IdToken": "mock_id_token",
            "RefreshToken": "mock_refresh_token",
        }
    }

    result = initiate_auth("test_user", "test_pass")
    assert result == {
        "AccessToken": "mock_access_token",
        "IdToken": "mock_id_token",
        "RefreshToken": "mock_refresh_token",
    }


def test_initiate_auth_with_secret_hash(mock_cognito_client):
    with patch(
        "app.auth.cognito.calculate_secret_hash", return_value="mocked_secret_hash"
    ) as mock_hash:
        mock_cognito_client.initiate_auth.return_value = {
            "AuthenticationResult": {"AccessToken": "token"}
        }

        initiate_auth("test_user", "test_pass")

        # Verify calculate_secret_hash was called
        mock_hash.assert_called_once_with(
            "test_user", TEST_CLIENT_ID, TEST_CLIENT_SECRET
        )

        # Verify SECRET_HASH was included in auth params
        call_args = mock_cognito_client.initiate_auth.call_args[1]
        assert "SECRET_HASH" in call_args["AuthParameters"]
        assert call_args["AuthParameters"]["SECRET_HASH"] == "mocked_secret_hash"


def test_initiate_auth_not_authorized(mock_cognito_client):
    mock_cognito_client.initiate_auth.side_effect = (
        mock_cognito_client.exceptions.NotAuthorizedException()
    )

    with pytest.raises(HTTPException) as exc_info:
        initiate_auth("invalid_user", "wrong_pass")

    assert exc_info.value.status_code == status.HTTP_401_UNAUTHORIZED
    assert exc_info.value.detail == "Invalid username or password"


def test_initiate_auth_user_not_found(mock_cognito_client):
    mock_cognito_client.initiate_auth.side_effect = (
        mock_cognito_client.exceptions.UserNotFoundException()
    )

    with pytest.raises(HTTPException) as exc_info:
        initiate_auth("nonexistent_user", "any_pass")

    assert exc_info.value.status_code == status.HTTP_404_NOT_FOUND
    assert exc_info.value.detail == "User not found"


def test_initiate_auth_generic_error(mock_cognito_client):
    mock_cognito_client.initiate_auth.side_effect = Exception("Some error")

    with pytest.raises(HTTPException) as exc_info:
        initiate_auth("test_user", "test_pass")

    assert exc_info.value.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
    assert "An error occurred during authentication" in exc_info.value.detail


def test_get_user_from_token_success(mock_cognito_client):
    mock_response = {
        "Username": "test_user",
        "UserAttributes": [
            {"Name": "sub", "Value": "123"},
            {"Name": USER_ROLE_ATTRIBUTE, "Value": "admin,user"},
        ],
    }
    mock_cognito_client.get_user.return_value = mock_response

    result = get_user_from_token("valid_token")

    assert isinstance(result, CognitoUser)
    assert result.username == "test_user"
    assert set(result.roles) == {"admin", "user"}


def test_get_user_from_token_no_roles(mock_cognito_client):
    mock_response = {
        "Username": "test_user",
        "UserAttributes": [{"Name": "sub", "Value": "123"}],
    }
    mock_cognito_client.get_user.return_value = mock_response

    result = get_user_from_token("valid_token")

    assert isinstance(result, CognitoUser)
    assert result.username == "test_user"
    assert result.roles == []


def test_get_user_from_token_invalid_token(mock_cognito_client):
    mock_cognito_client.get_user.side_effect = (
        mock_cognito_client.exceptions.NotAuthorizedException()
    )

    with pytest.raises(HTTPException) as exc_info:
        get_user_from_token("invalid_token")

    assert exc_info.value.status_code == status.HTTP_401_UNAUTHORIZED
    assert exc_info.value.detail == "Invalid or expired token."


def test_get_user_from_token_user_not_found(mock_cognito_client):
    mock_cognito_client.get_user.side_effect = (
        mock_cognito_client.exceptions.UserNotFoundException()
    )

    with pytest.raises(HTTPException) as exc_info:
        get_user_from_token("token_for_nonexistent_user")

    assert exc_info.value.status_code == status.HTTP_401_UNAUTHORIZED
    assert exc_info.value.detail == "User not found or invalid token."


def test_get_user_from_token_generic_error(mock_cognito_client):
    mock_cognito_client.get_user.side_effect = Exception("Some error")

    with pytest.raises(HTTPException) as exc_info:
        get_user_from_token("test_token")

    assert exc_info.value.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
    assert "Token verification failed" in exc_info.value.detail