mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-06 14:48:29 +00:00
fc3f2c49d2
* Remove the need of typing the email on 'crewai login' * Remove auth0 constants, update tests
339 lines
12 KiB
Python
339 lines
12 KiB
Python
import pytest
|
|
from datetime import datetime, timedelta
|
|
import requests
|
|
from unittest.mock import MagicMock, patch, call
|
|
from crewai.cli.authentication.main import AuthenticationCommand
|
|
from crewai.cli.constants import (
|
|
CREWAI_ENTERPRISE_DEFAULT_OAUTH2_DOMAIN,
|
|
CREWAI_ENTERPRISE_DEFAULT_OAUTH2_CLIENT_ID,
|
|
CREWAI_ENTERPRISE_DEFAULT_OAUTH2_AUDIENCE,
|
|
)
|
|
|
|
|
|
class TestAuthenticationCommand:
|
|
def setup_method(self):
|
|
self.auth_command = AuthenticationCommand()
|
|
|
|
@pytest.mark.parametrize(
|
|
"user_provider,expected_urls",
|
|
[
|
|
(
|
|
"workos",
|
|
{
|
|
"device_code_url": f"https://{CREWAI_ENTERPRISE_DEFAULT_OAUTH2_DOMAIN}/oauth2/device_authorization",
|
|
"token_url": f"https://{CREWAI_ENTERPRISE_DEFAULT_OAUTH2_DOMAIN}/oauth2/token",
|
|
"client_id": CREWAI_ENTERPRISE_DEFAULT_OAUTH2_CLIENT_ID,
|
|
"audience": CREWAI_ENTERPRISE_DEFAULT_OAUTH2_AUDIENCE,
|
|
"domain": CREWAI_ENTERPRISE_DEFAULT_OAUTH2_DOMAIN,
|
|
},
|
|
),
|
|
],
|
|
)
|
|
@patch("crewai.cli.authentication.main.AuthenticationCommand._get_device_code")
|
|
@patch(
|
|
"crewai.cli.authentication.main.AuthenticationCommand._display_auth_instructions"
|
|
)
|
|
@patch("crewai.cli.authentication.main.AuthenticationCommand._poll_for_token")
|
|
@patch("crewai.cli.authentication.main.console.print")
|
|
def test_login(
|
|
self,
|
|
mock_console_print,
|
|
mock_poll,
|
|
mock_display,
|
|
mock_get_device,
|
|
user_provider,
|
|
expected_urls,
|
|
):
|
|
mock_get_device.return_value = {
|
|
"device_code": "test_code",
|
|
"user_code": "123456",
|
|
}
|
|
|
|
self.auth_command.login()
|
|
|
|
mock_console_print.assert_called_once_with(
|
|
"Signing in to CrewAI Enterprise...\n", style="bold blue"
|
|
)
|
|
mock_get_device.assert_called_once()
|
|
mock_display.assert_called_once_with(
|
|
{"device_code": "test_code", "user_code": "123456"}
|
|
)
|
|
mock_poll.assert_called_once_with(
|
|
{"device_code": "test_code", "user_code": "123456"},
|
|
)
|
|
assert (
|
|
self.auth_command.oauth2_provider.get_client_id()
|
|
== expected_urls["client_id"]
|
|
)
|
|
assert (
|
|
self.auth_command.oauth2_provider.get_audience()
|
|
== expected_urls["audience"]
|
|
)
|
|
assert (
|
|
self.auth_command.oauth2_provider._get_domain() == expected_urls["domain"]
|
|
)
|
|
|
|
@patch("crewai.cli.authentication.main.webbrowser")
|
|
@patch("crewai.cli.authentication.main.console.print")
|
|
def test_display_auth_instructions(self, mock_console_print, mock_webbrowser):
|
|
device_code_data = {
|
|
"verification_uri_complete": "https://example.com/auth",
|
|
"user_code": "123456",
|
|
}
|
|
|
|
self.auth_command._display_auth_instructions(device_code_data)
|
|
|
|
expected_calls = [
|
|
call("1. Navigate to: ", "https://example.com/auth"),
|
|
call("2. Enter the following code: ", "123456"),
|
|
]
|
|
mock_console_print.assert_has_calls(expected_calls)
|
|
mock_webbrowser.open.assert_called_once_with("https://example.com/auth")
|
|
|
|
@pytest.mark.parametrize(
|
|
"user_provider,jwt_config",
|
|
[
|
|
(
|
|
"workos",
|
|
{
|
|
"jwks_url": f"https://{CREWAI_ENTERPRISE_DEFAULT_OAUTH2_DOMAIN}/oauth2/jwks",
|
|
"issuer": f"https://{CREWAI_ENTERPRISE_DEFAULT_OAUTH2_DOMAIN}",
|
|
"audience": CREWAI_ENTERPRISE_DEFAULT_OAUTH2_AUDIENCE,
|
|
},
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("has_expiration", [True, False])
|
|
@patch("crewai.cli.authentication.main.validate_jwt_token")
|
|
@patch("crewai.cli.authentication.main.TokenManager.save_tokens")
|
|
def test_validate_and_save_token(
|
|
self,
|
|
mock_save_tokens,
|
|
mock_validate_jwt,
|
|
user_provider,
|
|
jwt_config,
|
|
has_expiration,
|
|
):
|
|
from crewai.cli.authentication.providers.workos import WorkosProvider
|
|
from crewai.cli.authentication.main import Oauth2Settings
|
|
|
|
if user_provider == "workos":
|
|
self.auth_command.oauth2_provider = WorkosProvider(
|
|
settings=Oauth2Settings(
|
|
provider=user_provider,
|
|
client_id="test-client-id",
|
|
domain=CREWAI_ENTERPRISE_DEFAULT_OAUTH2_DOMAIN,
|
|
audience=jwt_config["audience"],
|
|
)
|
|
)
|
|
|
|
token_data = {"access_token": "test_access_token", "id_token": "test_id_token"}
|
|
|
|
if has_expiration:
|
|
future_timestamp = int((datetime.now() + timedelta(days=100)).timestamp())
|
|
decoded_token = {"exp": future_timestamp}
|
|
else:
|
|
decoded_token = {}
|
|
|
|
mock_validate_jwt.return_value = decoded_token
|
|
|
|
self.auth_command._validate_and_save_token(token_data)
|
|
|
|
mock_validate_jwt.assert_called_once_with(
|
|
jwt_token="test_access_token",
|
|
jwks_url=jwt_config["jwks_url"],
|
|
issuer=jwt_config["issuer"],
|
|
audience=jwt_config["audience"],
|
|
)
|
|
|
|
if has_expiration:
|
|
mock_save_tokens.assert_called_once_with(
|
|
"test_access_token", future_timestamp
|
|
)
|
|
else:
|
|
mock_save_tokens.assert_called_once_with("test_access_token", 0)
|
|
|
|
@patch("crewai.cli.tools.main.ToolCommand")
|
|
@patch("crewai.cli.authentication.main.Settings")
|
|
@patch("crewai.cli.authentication.main.console.print")
|
|
def test_login_to_tool_repository_success(
|
|
self, mock_console_print, mock_settings, mock_tool_command
|
|
):
|
|
mock_tool_instance = MagicMock()
|
|
mock_tool_command.return_value = mock_tool_instance
|
|
|
|
mock_settings_instance = MagicMock()
|
|
mock_settings_instance.org_name = "Test Org"
|
|
mock_settings_instance.org_uuid = "test-uuid-123"
|
|
mock_settings.return_value = mock_settings_instance
|
|
|
|
self.auth_command._login_to_tool_repository()
|
|
|
|
mock_tool_command.assert_called_once()
|
|
mock_tool_instance.login.assert_called_once()
|
|
|
|
expected_calls = [
|
|
call(
|
|
"Now logging you in to the Tool Repository... ",
|
|
style="bold blue",
|
|
end="",
|
|
),
|
|
call("Success!\n", style="bold green"),
|
|
call(
|
|
"You are authenticated to the tool repository as [bold cyan]'Test Org'[/bold cyan] (test-uuid-123)",
|
|
style="green",
|
|
),
|
|
]
|
|
mock_console_print.assert_has_calls(expected_calls)
|
|
|
|
@patch("crewai.cli.tools.main.ToolCommand")
|
|
@patch("crewai.cli.authentication.main.console.print")
|
|
def test_login_to_tool_repository_error(
|
|
self, mock_console_print, mock_tool_command
|
|
):
|
|
mock_tool_instance = MagicMock()
|
|
mock_tool_instance.login.side_effect = Exception("Tool repository error")
|
|
mock_tool_command.return_value = mock_tool_instance
|
|
|
|
self.auth_command._login_to_tool_repository()
|
|
|
|
mock_tool_command.assert_called_once()
|
|
mock_tool_instance.login.assert_called_once()
|
|
|
|
expected_calls = [
|
|
call(
|
|
"Now logging you in to the Tool Repository... ",
|
|
style="bold blue",
|
|
end="",
|
|
),
|
|
call(
|
|
"\n[bold yellow]Warning:[/bold yellow] Authentication with the Tool Repository failed.",
|
|
style="yellow",
|
|
),
|
|
call(
|
|
"Other features will work normally, but you may experience limitations with downloading and publishing tools.\nRun [bold]crewai login[/bold] to try logging in again.\n",
|
|
style="yellow",
|
|
),
|
|
]
|
|
mock_console_print.assert_has_calls(expected_calls)
|
|
|
|
@patch("requests.post")
|
|
def test_get_device_code(self, mock_post):
|
|
mock_response = MagicMock()
|
|
mock_response.json.return_value = {
|
|
"device_code": "test_device_code",
|
|
"user_code": "123456",
|
|
"verification_uri_complete": "https://example.com/auth",
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
self.auth_command.oauth2_provider = MagicMock()
|
|
self.auth_command.oauth2_provider.get_client_id.return_value = "test_client"
|
|
self.auth_command.oauth2_provider.get_authorize_url.return_value = (
|
|
"https://example.com/device"
|
|
)
|
|
self.auth_command.oauth2_provider.get_audience.return_value = "test_audience"
|
|
|
|
result = self.auth_command._get_device_code()
|
|
|
|
mock_post.assert_called_once_with(
|
|
url="https://example.com/device",
|
|
data={
|
|
"client_id": "test_client",
|
|
"scope": "openid",
|
|
"audience": "test_audience",
|
|
},
|
|
timeout=20,
|
|
)
|
|
|
|
assert result == {
|
|
"device_code": "test_device_code",
|
|
"user_code": "123456",
|
|
"verification_uri_complete": "https://example.com/auth",
|
|
}
|
|
|
|
@patch("requests.post")
|
|
@patch("crewai.cli.authentication.main.console.print")
|
|
def test_poll_for_token_success(self, mock_console_print, mock_post):
|
|
mock_response_success = MagicMock()
|
|
mock_response_success.status_code = 200
|
|
mock_response_success.json.return_value = {
|
|
"access_token": "test_access_token",
|
|
"id_token": "test_id_token",
|
|
}
|
|
mock_post.return_value = mock_response_success
|
|
|
|
device_code_data = {"device_code": "test_device_code", "interval": 1}
|
|
|
|
with (
|
|
patch.object(
|
|
self.auth_command, "_validate_and_save_token"
|
|
) as mock_validate,
|
|
patch.object(
|
|
self.auth_command, "_login_to_tool_repository"
|
|
) as mock_tool_login,
|
|
):
|
|
self.auth_command.oauth2_provider = MagicMock()
|
|
self.auth_command.oauth2_provider.get_token_url.return_value = (
|
|
"https://example.com/token"
|
|
)
|
|
self.auth_command.oauth2_provider.get_client_id.return_value = "test_client"
|
|
|
|
self.auth_command._poll_for_token(device_code_data)
|
|
|
|
mock_post.assert_called_once_with(
|
|
"https://example.com/token",
|
|
data={
|
|
"grant_type": "urn:ietf:params:oauth:grant-type:device_code",
|
|
"device_code": "test_device_code",
|
|
"client_id": "test_client",
|
|
},
|
|
timeout=30,
|
|
)
|
|
|
|
mock_validate.assert_called_once()
|
|
mock_tool_login.assert_called_once()
|
|
|
|
expected_calls = [
|
|
call("\nWaiting for authentication... ", style="bold blue", end=""),
|
|
call("Success!", style="bold green"),
|
|
call("\n[bold green]Welcome to CrewAI Enterprise![/bold green]\n"),
|
|
]
|
|
mock_console_print.assert_has_calls(expected_calls)
|
|
|
|
@patch("requests.post")
|
|
@patch("crewai.cli.authentication.main.console.print")
|
|
def test_poll_for_token_timeout(self, mock_console_print, mock_post):
|
|
mock_response_pending = MagicMock()
|
|
mock_response_pending.status_code = 400
|
|
mock_response_pending.json.return_value = {"error": "authorization_pending"}
|
|
mock_post.return_value = mock_response_pending
|
|
|
|
device_code_data = {
|
|
"device_code": "test_device_code",
|
|
"interval": 0.1, # Short interval for testing
|
|
}
|
|
|
|
self.auth_command._poll_for_token(device_code_data)
|
|
|
|
mock_console_print.assert_any_call(
|
|
"Timeout: Failed to get the token. Please try again.", style="bold red"
|
|
)
|
|
|
|
@patch("requests.post")
|
|
def test_poll_for_token_error(self, mock_post):
|
|
"""Test the method to poll for token (error path)."""
|
|
# Setup mock to return error
|
|
mock_response_error = MagicMock()
|
|
mock_response_error.status_code = 400
|
|
mock_response_error.json.return_value = {
|
|
"error": "access_denied",
|
|
"error_description": "User denied access",
|
|
}
|
|
mock_post.return_value = mock_response_error
|
|
|
|
device_code_data = {"device_code": "test_device_code", "interval": 1}
|
|
|
|
with pytest.raises(requests.HTTPError):
|
|
self.auth_command._poll_for_token(device_code_data)
|