Added unit tests for cognito.py

tests/auth/test_cognito.py

@@ -0,0 +1,143 @@
+import pytest
+from unittest.mock import patch, MagicMock
+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 initiate_auth, get_user_from_token
+    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"}
+        }
+        
+        result = 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