mirror of
https://github.com/Aider-AI/aider.git
synced 2025-05-29 16:54:59 +00:00
style: Apply linter fixes to onboarding tests
This commit is contained in:
Paul Gauthier (aider)
parent
b54629addb
commit
24c074eeaa
1 changed files with 134 additions and 96 deletions
|
|
@ -1,41 +1,54 @@
|
|||
import unittest
|
||||
from unittest.mock import MagicMock, patch, mock_open
|
||||
import os
|
||||
import requests
|
||||
import socketserver
|
||||
import secrets
|
||||
import hashlib
|
||||
import base64
|
||||
import argparse
|
||||
import base64
|
||||
import hashlib
|
||||
import os
|
||||
import secrets
|
||||
import socketserver
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, mock_open, patch
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
# Mock the Analytics class as it's used in some functions
|
||||
class DummyAnalytics:
|
||||
def event(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
# Mock the InputOutput class
|
||||
class DummyIO:
|
||||
def tool_output(self, *args, **kwargs): pass
|
||||
def tool_warning(self, *args, **kwargs): pass
|
||||
def tool_error(self, *args, **kwargs): pass
|
||||
def confirm_ask(self, *args, **kwargs): return False # Default to no confirmation
|
||||
def offer_url(self, *args, **kwargs): pass
|
||||
def tool_output(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def tool_warning(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def tool_error(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def confirm_ask(self, *args, **kwargs):
|
||||
return False # Default to no confirmation
|
||||
|
||||
def offer_url(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
# Import the functions to be tested
|
||||
from aider.onboarding import (
|
||||
check_openrouter_tier,
|
||||
try_to_select_default_model,
|
||||
select_default_model,
|
||||
offer_openrouter_oauth,
|
||||
exchange_code_for_key,
|
||||
find_available_port,
|
||||
generate_pkce_codes,
|
||||
exchange_code_for_key,
|
||||
offer_openrouter_oauth,
|
||||
select_default_model,
|
||||
start_openrouter_oauth_flow,
|
||||
try_to_select_default_model,
|
||||
)
|
||||
|
||||
class TestOnboarding(unittest.TestCase):
|
||||
|
||||
@patch('requests.get')
|
||||
class TestOnboarding(unittest.TestCase):
|
||||
@patch("requests.get")
|
||||
def test_check_openrouter_tier_free(self, mock_get):
|
||||
"""Test check_openrouter_tier identifies free tier."""
|
||||
mock_response = MagicMock()
|
||||
|
|
@ -49,7 +62,7 @@ class TestOnboarding(unittest.TestCase):
|
|||
timeout=5,
|
||||
)
|
||||
|
||||
@patch('requests.get')
|
||||
@patch("requests.get")
|
||||
def test_check_openrouter_tier_paid(self, mock_get):
|
||||
"""Test check_openrouter_tier identifies paid tier."""
|
||||
mock_response = MagicMock()
|
||||
|
|
@ -58,13 +71,13 @@ class TestOnboarding(unittest.TestCase):
|
|||
mock_get.return_value = mock_response
|
||||
self.assertFalse(check_openrouter_tier("fake_key"))
|
||||
|
||||
@patch('requests.get')
|
||||
@patch("requests.get")
|
||||
def test_check_openrouter_tier_api_error(self, mock_get):
|
||||
"""Test check_openrouter_tier defaults to free on API error."""
|
||||
mock_get.side_effect = requests.exceptions.RequestException("API Error")
|
||||
self.assertTrue(check_openrouter_tier("fake_key"))
|
||||
|
||||
@patch('requests.get')
|
||||
@patch("requests.get")
|
||||
def test_check_openrouter_tier_missing_key(self, mock_get):
|
||||
"""Test check_openrouter_tier defaults to free if key is missing in response."""
|
||||
mock_response = MagicMock()
|
||||
|
|
@ -73,77 +86,81 @@ class TestOnboarding(unittest.TestCase):
|
|||
mock_get.return_value = mock_response
|
||||
self.assertTrue(check_openrouter_tier("fake_key"))
|
||||
|
||||
@patch('aider.onboarding.check_openrouter_tier')
|
||||
@patch("aider.onboarding.check_openrouter_tier")
|
||||
@patch.dict(os.environ, {}, clear=True)
|
||||
def test_try_select_default_model_no_keys(self, mock_check_tier):
|
||||
"""Test no model is selected when no keys are present."""
|
||||
self.assertIsNone(try_to_select_default_model())
|
||||
mock_check_tier.assert_not_called()
|
||||
|
||||
@patch('aider.onboarding.check_openrouter_tier', return_value=True) # Assume free tier
|
||||
@patch("aider.onboarding.check_openrouter_tier", return_value=True) # Assume free tier
|
||||
@patch.dict(os.environ, {"OPENROUTER_API_KEY": "or_key"}, clear=True)
|
||||
def test_try_select_default_model_openrouter_free(self, mock_check_tier):
|
||||
"""Test OpenRouter free model selection."""
|
||||
self.assertEqual(try_to_select_default_model(), "openrouter/google/gemini-2.5-pro-exp-03-25:free")
|
||||
self.assertEqual(
|
||||
try_to_select_default_model(), "openrouter/google/gemini-2.5-pro-exp-03-25:free"
|
||||
)
|
||||
mock_check_tier.assert_called_once_with("or_key")
|
||||
|
||||
@patch('aider.onboarding.check_openrouter_tier', return_value=False) # Assume paid tier
|
||||
@patch("aider.onboarding.check_openrouter_tier", return_value=False) # Assume paid tier
|
||||
@patch.dict(os.environ, {"OPENROUTER_API_KEY": "or_key"}, clear=True)
|
||||
def test_try_select_default_model_openrouter_paid(self, mock_check_tier):
|
||||
"""Test OpenRouter paid model selection."""
|
||||
self.assertEqual(try_to_select_default_model(), "openrouter/anthropic/claude-3.7-sonnet")
|
||||
mock_check_tier.assert_called_once_with("or_key")
|
||||
|
||||
@patch('aider.onboarding.check_openrouter_tier')
|
||||
@patch("aider.onboarding.check_openrouter_tier")
|
||||
@patch.dict(os.environ, {"ANTHROPIC_API_KEY": "an_key"}, clear=True)
|
||||
def test_try_select_default_model_anthropic(self, mock_check_tier):
|
||||
"""Test Anthropic model selection."""
|
||||
self.assertEqual(try_to_select_default_model(), "sonnet")
|
||||
mock_check_tier.assert_not_called()
|
||||
|
||||
@patch('aider.onboarding.check_openrouter_tier')
|
||||
@patch("aider.onboarding.check_openrouter_tier")
|
||||
@patch.dict(os.environ, {"DEEPSEEK_API_KEY": "ds_key"}, clear=True)
|
||||
def test_try_select_default_model_deepseek(self, mock_check_tier):
|
||||
"""Test Deepseek model selection."""
|
||||
self.assertEqual(try_to_select_default_model(), "deepseek")
|
||||
mock_check_tier.assert_not_called()
|
||||
|
||||
@patch('aider.onboarding.check_openrouter_tier')
|
||||
@patch("aider.onboarding.check_openrouter_tier")
|
||||
@patch.dict(os.environ, {"OPENAI_API_KEY": "oa_key"}, clear=True)
|
||||
def test_try_select_default_model_openai(self, mock_check_tier):
|
||||
"""Test OpenAI model selection."""
|
||||
self.assertEqual(try_to_select_default_model(), "gpt-4o")
|
||||
mock_check_tier.assert_not_called()
|
||||
|
||||
@patch('aider.onboarding.check_openrouter_tier')
|
||||
@patch("aider.onboarding.check_openrouter_tier")
|
||||
@patch.dict(os.environ, {"GEMINI_API_KEY": "gm_key"}, clear=True)
|
||||
def test_try_select_default_model_gemini(self, mock_check_tier):
|
||||
"""Test Gemini model selection."""
|
||||
self.assertEqual(try_to_select_default_model(), "gemini/gemini-2.5-pro-exp-03-25")
|
||||
mock_check_tier.assert_not_called()
|
||||
|
||||
@patch('aider.onboarding.check_openrouter_tier')
|
||||
@patch("aider.onboarding.check_openrouter_tier")
|
||||
@patch.dict(os.environ, {"VERTEXAI_PROJECT": "vx_proj"}, clear=True)
|
||||
def test_try_select_default_model_vertex(self, mock_check_tier):
|
||||
"""Test Vertex AI model selection."""
|
||||
self.assertEqual(try_to_select_default_model(), "vertex_ai/gemini-2.5-pro-exp-03-25")
|
||||
mock_check_tier.assert_not_called()
|
||||
|
||||
@patch('aider.onboarding.check_openrouter_tier', return_value=False) # Paid
|
||||
@patch.dict(os.environ, {"OPENROUTER_API_KEY": "or_key", "OPENAI_API_KEY": "oa_key"}, clear=True)
|
||||
@patch("aider.onboarding.check_openrouter_tier", return_value=False) # Paid
|
||||
@patch.dict(
|
||||
os.environ, {"OPENROUTER_API_KEY": "or_key", "OPENAI_API_KEY": "oa_key"}, clear=True
|
||||
)
|
||||
def test_try_select_default_model_priority_openrouter(self, mock_check_tier):
|
||||
"""Test OpenRouter key takes priority."""
|
||||
self.assertEqual(try_to_select_default_model(), "openrouter/anthropic/claude-3.7-sonnet")
|
||||
mock_check_tier.assert_called_once_with("or_key")
|
||||
|
||||
@patch('aider.onboarding.check_openrouter_tier')
|
||||
@patch("aider.onboarding.check_openrouter_tier")
|
||||
@patch.dict(os.environ, {"ANTHROPIC_API_KEY": "an_key", "OPENAI_API_KEY": "oa_key"}, clear=True)
|
||||
def test_try_select_default_model_priority_anthropic(self, mock_check_tier):
|
||||
"""Test Anthropic key takes priority over OpenAI."""
|
||||
self.assertEqual(try_to_select_default_model(), "sonnet")
|
||||
mock_check_tier.assert_not_called()
|
||||
|
||||
@patch('socketserver.TCPServer')
|
||||
@patch("socketserver.TCPServer")
|
||||
def test_find_available_port_success(self, mock_tcp_server):
|
||||
"""Test finding an available port."""
|
||||
# Simulate port 8484 being available
|
||||
|
|
@ -152,7 +169,7 @@ class TestOnboarding(unittest.TestCase):
|
|||
self.assertEqual(port, 8484)
|
||||
mock_tcp_server.assert_called_once_with(("localhost", 8484), None)
|
||||
|
||||
@patch('socketserver.TCPServer')
|
||||
@patch("socketserver.TCPServer")
|
||||
def test_find_available_port_in_use(self, mock_tcp_server):
|
||||
"""Test finding the next available port if the first is in use."""
|
||||
# Simulate port 8484 raising OSError, 8485 being available
|
||||
|
|
@ -164,7 +181,7 @@ class TestOnboarding(unittest.TestCase):
|
|||
mock_tcp_server.assert_any_call(("localhost", 8484), None)
|
||||
mock_tcp_server.assert_any_call(("localhost", 8485), None)
|
||||
|
||||
@patch('socketserver.TCPServer', side_effect=OSError)
|
||||
@patch("socketserver.TCPServer", side_effect=OSError)
|
||||
def test_find_available_port_none_available(self, mock_tcp_server):
|
||||
"""Test returning None if no ports are available in the range."""
|
||||
port = find_available_port(start_port=8484, end_port=8485)
|
||||
|
|
@ -184,7 +201,7 @@ class TestOnboarding(unittest.TestCase):
|
|||
expected_challenge = base64.urlsafe_b64encode(hasher.digest()).rstrip(b"=").decode("utf-8")
|
||||
self.assertEqual(challenge, expected_challenge)
|
||||
|
||||
@patch('requests.post')
|
||||
@patch("requests.post")
|
||||
def test_exchange_code_for_key_success(self, mock_post):
|
||||
"""Test successful code exchange for API key."""
|
||||
mock_response = MagicMock()
|
||||
|
|
@ -207,7 +224,7 @@ class TestOnboarding(unittest.TestCase):
|
|||
timeout=30,
|
||||
)
|
||||
|
||||
@patch('requests.post')
|
||||
@patch("requests.post")
|
||||
def test_exchange_code_for_key_missing_key(self, mock_post):
|
||||
"""Test code exchange when 'key' is missing in response."""
|
||||
mock_response = MagicMock()
|
||||
|
|
@ -224,7 +241,7 @@ class TestOnboarding(unittest.TestCase):
|
|||
io_mock.tool_error.assert_any_call("Error: 'key' not found in OpenRouter response.")
|
||||
io_mock.tool_error.assert_any_call('Response: {"other_data": "value"}')
|
||||
|
||||
@patch('requests.post')
|
||||
@patch("requests.post")
|
||||
def test_exchange_code_for_key_http_error(self, mock_post):
|
||||
"""Test code exchange with HTTP error."""
|
||||
mock_response = MagicMock()
|
||||
|
|
@ -244,7 +261,7 @@ class TestOnboarding(unittest.TestCase):
|
|||
)
|
||||
io_mock.tool_error.assert_any_call('Response: {"error": "invalid_code"}')
|
||||
|
||||
@patch('requests.post')
|
||||
@patch("requests.post")
|
||||
def test_exchange_code_for_key_timeout(self, mock_post):
|
||||
"""Test code exchange with timeout."""
|
||||
mock_post.side_effect = requests.exceptions.Timeout("Timeout")
|
||||
|
|
@ -258,7 +275,7 @@ class TestOnboarding(unittest.TestCase):
|
|||
"Error: Request to OpenRouter timed out during code exchange."
|
||||
)
|
||||
|
||||
@patch('requests.post')
|
||||
@patch("requests.post")
|
||||
def test_exchange_code_for_key_request_exception(self, mock_post):
|
||||
"""Test code exchange with general request exception."""
|
||||
req_exception = requests.exceptions.RequestException("Network Error")
|
||||
|
|
@ -275,8 +292,8 @@ class TestOnboarding(unittest.TestCase):
|
|||
|
||||
# --- Tests for select_default_model ---
|
||||
|
||||
@patch('aider.onboarding.try_to_select_default_model', return_value="gpt-4o")
|
||||
@patch('aider.onboarding.offer_openrouter_oauth')
|
||||
@patch("aider.onboarding.try_to_select_default_model", return_value="gpt-4o")
|
||||
@patch("aider.onboarding.offer_openrouter_oauth")
|
||||
def test_select_default_model_already_specified(self, mock_offer_oauth, mock_try_select):
|
||||
"""Test select_default_model returns args.model if provided."""
|
||||
args = argparse.Namespace(model="specific-model")
|
||||
|
|
@ -287,8 +304,8 @@ class TestOnboarding(unittest.TestCase):
|
|||
mock_try_select.assert_not_called()
|
||||
mock_offer_oauth.assert_not_called()
|
||||
|
||||
@patch('aider.onboarding.try_to_select_default_model', return_value="gpt-4o")
|
||||
@patch('aider.onboarding.offer_openrouter_oauth')
|
||||
@patch("aider.onboarding.try_to_select_default_model", return_value="gpt-4o")
|
||||
@patch("aider.onboarding.offer_openrouter_oauth")
|
||||
def test_select_default_model_found_via_env(self, mock_offer_oauth, mock_try_select):
|
||||
"""Test select_default_model returns model found by try_to_select."""
|
||||
args = argparse.Namespace(model=None) # No model specified
|
||||
|
|
@ -307,8 +324,12 @@ class TestOnboarding(unittest.TestCase):
|
|||
analytics_mock.event.assert_called_once_with("auto_model_selection", model="gpt-4o")
|
||||
mock_offer_oauth.assert_not_called()
|
||||
|
||||
@patch('aider.onboarding.try_to_select_default_model', side_effect=[None, None]) # Fails first, fails after oauth attempt
|
||||
@patch('aider.onboarding.offer_openrouter_oauth', return_value=False) # OAuth offered but fails/declined
|
||||
@patch(
|
||||
"aider.onboarding.try_to_select_default_model", side_effect=[None, None]
|
||||
) # Fails first, fails after oauth attempt
|
||||
@patch(
|
||||
"aider.onboarding.offer_openrouter_oauth", return_value=False
|
||||
) # OAuth offered but fails/declined
|
||||
def test_select_default_model_no_keys_oauth_fail(self, mock_offer_oauth, mock_try_select):
|
||||
"""Test select_default_model offers OAuth when no keys, but OAuth fails."""
|
||||
args = argparse.Namespace(model=None)
|
||||
|
|
@ -327,8 +348,13 @@ class TestOnboarding(unittest.TestCase):
|
|||
)
|
||||
io_mock.offer_url.assert_called_once() # Should offer docs URL
|
||||
|
||||
@patch('aider.onboarding.try_to_select_default_model', side_effect=[None, "openrouter/google/gemini-2.5-pro-exp-03-25:free"]) # Fails first, succeeds after oauth
|
||||
@patch('aider.onboarding.offer_openrouter_oauth', return_value=True) # OAuth offered and succeeds
|
||||
@patch(
|
||||
"aider.onboarding.try_to_select_default_model",
|
||||
side_effect=[None, "openrouter/google/gemini-2.5-pro-exp-03-25:free"],
|
||||
) # Fails first, succeeds after oauth
|
||||
@patch(
|
||||
"aider.onboarding.offer_openrouter_oauth", return_value=True
|
||||
) # OAuth offered and succeeds
|
||||
def test_select_default_model_no_keys_oauth_success(self, mock_offer_oauth, mock_try_select):
|
||||
"""Test select_default_model offers OAuth, which succeeds."""
|
||||
args = argparse.Namespace(model=None)
|
||||
|
|
@ -353,7 +379,7 @@ class TestOnboarding(unittest.TestCase):
|
|||
# Instead, let's verify the final state and model returned.
|
||||
|
||||
# --- Tests for offer_openrouter_oauth ---
|
||||
@patch('aider.onboarding.start_openrouter_oauth_flow', return_value="new_or_key")
|
||||
@patch("aider.onboarding.start_openrouter_oauth_flow", return_value="new_or_key")
|
||||
@patch.dict(os.environ, {}, clear=True) # Ensure no key exists initially
|
||||
def test_offer_openrouter_oauth_confirm_yes_success(self, mock_start_oauth):
|
||||
"""Test offer_openrouter_oauth when user confirms and OAuth succeeds."""
|
||||
|
|
@ -373,7 +399,7 @@ class TestOnboarding(unittest.TestCase):
|
|||
# Clean up env var
|
||||
del os.environ["OPENROUTER_API_KEY"]
|
||||
|
||||
@patch('aider.onboarding.start_openrouter_oauth_flow', return_value=None) # OAuth fails
|
||||
@patch("aider.onboarding.start_openrouter_oauth_flow", return_value=None) # OAuth fails
|
||||
@patch.dict(os.environ, {}, clear=True)
|
||||
def test_offer_openrouter_oauth_confirm_yes_fail(self, mock_start_oauth):
|
||||
"""Test offer_openrouter_oauth when user confirms but OAuth fails."""
|
||||
|
|
@ -389,12 +415,13 @@ class TestOnboarding(unittest.TestCase):
|
|||
io_mock.confirm_ask.assert_called_once()
|
||||
mock_start_oauth.assert_called_once_with(io_mock, analytics_mock)
|
||||
self.assertNotIn("OPENROUTER_API_KEY", os.environ)
|
||||
io_mock.tool_error.assert_called_once_with("OpenRouter authentication did not complete successfully.")
|
||||
io_mock.tool_error.assert_called_once_with(
|
||||
"OpenRouter authentication did not complete successfully."
|
||||
)
|
||||
analytics_mock.event.assert_any_call("oauth_flow_initiated", provider="openrouter")
|
||||
analytics_mock.event.assert_any_call("oauth_flow_failure")
|
||||
|
||||
|
||||
@patch('aider.onboarding.start_openrouter_oauth_flow')
|
||||
@patch("aider.onboarding.start_openrouter_oauth_flow")
|
||||
def test_offer_openrouter_oauth_confirm_no(self, mock_start_oauth):
|
||||
"""Test offer_openrouter_oauth when user declines."""
|
||||
io_mock = DummyIO()
|
||||
|
|
@ -409,21 +436,30 @@ class TestOnboarding(unittest.TestCase):
|
|||
mock_start_oauth.assert_not_called()
|
||||
analytics_mock.event.assert_not_called() # No OAuth events if declined
|
||||
|
||||
|
||||
# --- More complex test for start_openrouter_oauth_flow (simplified) ---
|
||||
# This test focuses on the successful path, mocking heavily
|
||||
@patch('aider.onboarding.check_pip_install_extra', return_value=True) # Assume requests is installed
|
||||
@patch('aider.onboarding.find_available_port', return_value=8484)
|
||||
@patch('threading.Thread')
|
||||
@patch('threading.Event')
|
||||
@patch('webbrowser.open')
|
||||
@patch('aider.onboarding.exchange_code_for_key', return_value="oauth_api_key")
|
||||
@patch('os.makedirs')
|
||||
@patch('builtins.open', new_callable=mock_open)
|
||||
@patch(
|
||||
"aider.onboarding.check_pip_install_extra", return_value=True
|
||||
) # Assume requests is installed
|
||||
@patch("aider.onboarding.find_available_port", return_value=8484)
|
||||
@patch("threading.Thread")
|
||||
@patch("threading.Event")
|
||||
@patch("webbrowser.open")
|
||||
@patch("aider.onboarding.exchange_code_for_key", return_value="oauth_api_key")
|
||||
@patch("os.makedirs")
|
||||
@patch("builtins.open", new_callable=mock_open)
|
||||
@patch.dict(os.environ, {}, clear=True) # Start with clean env
|
||||
def test_start_openrouter_oauth_flow_success_path(
|
||||
self, mock_env, mock_open_file, mock_makedirs, mock_exchange, mock_webbrowser,
|
||||
mock_event_cls, mock_thread_cls, mock_find_port, mock_check_pip
|
||||
self,
|
||||
mock_env,
|
||||
mock_open_file,
|
||||
mock_makedirs,
|
||||
mock_exchange,
|
||||
mock_webbrowser,
|
||||
mock_event_cls,
|
||||
mock_thread_cls,
|
||||
mock_find_port,
|
||||
mock_check_pip,
|
||||
):
|
||||
"""Test the successful path of start_openrouter_oauth_flow."""
|
||||
io_mock = DummyIO()
|
||||
|
|
@ -469,7 +505,9 @@ class TestOnboarding(unittest.TestCase):
|
|||
self.assertEqual(os.environ.get("OPENROUTER_API_KEY"), "oauth_api_key")
|
||||
# Check that saving the key was attempted
|
||||
mock_makedirs.assert_called_once()
|
||||
mock_open_file.assert_called_once_with(os.path.expanduser("~/.aider/oauth-keys.env"), "a", encoding="utf-8")
|
||||
mock_open_file.assert_called_once_with(
|
||||
os.path.expanduser("~/.aider/oauth-keys.env"), "a", encoding="utf-8"
|
||||
)
|
||||
mock_open_file().write.assert_called_once_with('OPENROUTER_API_KEY="oauth_api_key"\n')
|
||||
# Check analytics events
|
||||
analytics_mock.event.assert_any_call("oauth_flow_code_received", provider="openrouter")
|
||||
|
|
@ -478,5 +516,5 @@ class TestOnboarding(unittest.TestCase):
|
|||
del os.environ["OPENROUTER_API_KEY"]
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue