mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-11 09:08:31 +00:00
103 lines
3.9 KiB
Python
103 lines
3.9 KiB
Python
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
import requests
|
|
|
|
from crewai.cli.authentication.main import AuthenticationCommand
|
|
|
|
|
|
class TestAuthenticationCommand(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.auth_command = AuthenticationCommand()
|
|
|
|
@patch("crewai.cli.authentication.main.requests.post")
|
|
def test_get_device_code(self, mock_post) -> None:
|
|
mock_response = MagicMock()
|
|
mock_response.json.return_value = {
|
|
"device_code": "123456",
|
|
"user_code": "ABCDEF",
|
|
"verification_uri_complete": "https://example.com",
|
|
"interval": 5,
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
device_code_data = self.auth_command._get_device_code()
|
|
|
|
assert device_code_data["device_code"] == "123456"
|
|
assert device_code_data["user_code"] == "ABCDEF"
|
|
assert device_code_data["verification_uri_complete"] == "https://example.com"
|
|
assert device_code_data["interval"] == 5
|
|
|
|
@patch("crewai.cli.authentication.main.console.print")
|
|
@patch("crewai.cli.authentication.main.webbrowser.open")
|
|
def test_display_auth_instructions(self, mock_open, mock_print) -> None:
|
|
device_code_data = {
|
|
"verification_uri_complete": "https://example.com",
|
|
"user_code": "ABCDEF",
|
|
}
|
|
|
|
self.auth_command._display_auth_instructions(device_code_data)
|
|
|
|
mock_print.assert_any_call("1. Navigate to: ", "https://example.com")
|
|
mock_print.assert_any_call("2. Enter the following code: ", "ABCDEF")
|
|
mock_open.assert_called_once_with("https://example.com")
|
|
|
|
@patch("crewai.cli.authentication.main.ToolCommand")
|
|
@patch("crewai.cli.authentication.main.requests.post")
|
|
@patch("crewai.cli.authentication.main.validate_token")
|
|
@patch("crewai.cli.authentication.main.console.print")
|
|
def test_poll_for_token_success(
|
|
self, mock_print, mock_validate_token, mock_post, mock_tool,
|
|
) -> None:
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"id_token": "TOKEN",
|
|
"access_token": "ACCESS_TOKEN",
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
mock_instance = mock_tool.return_value
|
|
mock_instance.login.return_value = None
|
|
|
|
self.auth_command._poll_for_token({"device_code": "123456"})
|
|
|
|
mock_validate_token.assert_called_once_with("TOKEN")
|
|
mock_print.assert_called_once_with(
|
|
"\n[bold green]Welcome to CrewAI Enterprise![/bold green]\n",
|
|
)
|
|
|
|
@patch("crewai.cli.authentication.main.requests.post")
|
|
@patch("crewai.cli.authentication.main.console.print")
|
|
def test_poll_for_token_error(self, mock_print, mock_post) -> None:
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 400
|
|
mock_response.json.return_value = {
|
|
"error": "invalid_request",
|
|
"error_description": "Invalid request",
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
with pytest.raises(requests.HTTPError):
|
|
self.auth_command._poll_for_token({"device_code": "123456"})
|
|
|
|
mock_print.assert_not_called()
|
|
|
|
@patch("crewai.cli.authentication.main.requests.post")
|
|
@patch("crewai.cli.authentication.main.console.print")
|
|
def test_poll_for_token_timeout(self, mock_print, mock_post) -> None:
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 400
|
|
mock_response.json.return_value = {
|
|
"error": "authorization_pending",
|
|
"error_description": "Authorization pending",
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
self.auth_command._poll_for_token({"device_code": "123456", "interval": 0.01})
|
|
|
|
mock_print.assert_called_once_with(
|
|
"Timeout: Failed to get the token. Please try again.", style="bold red",
|
|
)
|