tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

feat: saved api_key to keychain for user (#104)

Browse Source
This commit is contained in:
Lifei Zhou
2024-10-07 09:40:50 +11:00
committed by GitHub
parent dfecf829a8
commit f3aab127e5
6 changed files with 238 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/cli/test_session.py
+1 -34
View File
@@ -2,8 +2,6 @@ from unittest.mock import MagicMock, patch
import pytest
from exchange import Exchange, Message, ToolUse, ToolResult
from exchange.providers.base import MissingProviderEnvVariableError
from exchange.invalid_choice_error import InvalidChoiceError
from goose.cli.prompt.goose_prompt_session import GoosePromptSession
from goose.cli.prompt.user_input import PromptAction, UserInput
from goose.cli.session import Session
@@ -22,7 +20,7 @@ def mock_specified_session_name():
@pytest.fixture
def create_session_with_mock_configs(mock_sessions_path, exchange_factory, profile_factory):
with (
patch("goose.cli.session.build_exchange") as mock_exchange,
patch("goose.cli.session.create_exchange") as mock_exchange,
patch("goose.cli.session.load_profile", return_value=profile_factory()),
patch("goose.cli.session.SessionNotifier") as mock_session_notifier,
patch("goose.cli.session.load_provider", return_value="provider"),
@@ -158,34 +156,3 @@ def test_set_generated_session_name(create_session_with_mock_configs, mock_sessi
with patch("goose.cli.session.droid", return_value=generated_session_name):
session = create_session_with_mock_configs({"name": None})
assert session.name == generated_session_name
def test_create_exchange_exit_when_env_var_does_not_exist(create_session_with_mock_configs, mock_sessions_path):
session = create_session_with_mock_configs()
expected_error = MissingProviderEnvVariableError(env_variable="OPENAI_API_KEY", provider="openai")
with (
patch("goose.cli.session.build_exchange", side_effect=expected_error),
patch("goose.cli.session.print") as mock_print,
patch("sys.exit") as mock_exit,
):
session._create_exchange()
mock_print.call_args_list[0][0][0].renderable == (
"Missing environment variable OPENAI_API_KEY for provider openai. ",
"Please set the required environment variable to continue.",
)
mock_exit.assert_called_once_with(1)
def test_create_exchange_exit_when_configuration_is_incorrect(create_session_with_mock_configs, mock_sessions_path):
session = create_session_with_mock_configs()
expected_error = InvalidChoiceError(
attribute_name="provider", attribute_value="wrong_provider", available_values=["openai"]
)
with (
patch("goose.cli.session.build_exchange", side_effect=expected_error),
patch("goose.cli.session.print") as mock_print,
patch("sys.exit") as mock_exit,
):
session._create_exchange()
assert "Unknown provider: wrong_provider. Available providers: openai" in mock_print.call_args_list[0][0][0]
mock_exit.assert_called_once_with(1)
tests/utils/test_create_exchange.py
+151
View File
@@ -0,0 +1,151 @@
import os
from unittest.mock import MagicMock, patch
from exchange.exchange import Exchange
from exchange.invalid_choice_error import InvalidChoiceError
from exchange.providers.base import MissingProviderEnvVariableError
import pytest
from goose.notifier import Notifier
from goose.profile import Profile
from goose.utils._create_exchange import create_exchange
TEST_PROFILE = MagicMock(spec=Profile)
TEST_EXCHANGE = MagicMock(spec=Exchange)
TEST_NOTIFIER = MagicMock(spec=Notifier)
@pytest.fixture
def mock_print():
with patch("goose.utils._create_exchange.print") as mock_print:
yield mock_print
@pytest.fixture
def mock_prompt():
with patch("goose.utils._create_exchange.prompt") as mock_prompt:
yield mock_prompt
@pytest.fixture
def mock_confirm():
with patch("goose.utils._create_exchange.confirm") as mock_confirm:
yield mock_confirm
@pytest.fixture
def mock_sys_exit():
with patch("sys.exit") as mock_exit:
yield mock_exit
@pytest.fixture
def mock_keyring_get_password():
with patch("keyring.get_password") as mock_get_password:
yield mock_get_password
@pytest.fixture
def mock_keyring_set_password():
with patch("keyring.set_password") as mock_set_password:
yield mock_set_password
def test_create_exchange_success(mock_print):
with patch("goose.utils._create_exchange.build_exchange", return_value=TEST_EXCHANGE):
assert create_exchange(profile=TEST_PROFILE, notifier=TEST_NOTIFIER) == TEST_EXCHANGE
def test_create_exchange_fail_with_invalid_choice_error(mock_print, mock_sys_exit):
expected_error = InvalidChoiceError(
attribute_name="provider", attribute_value="wrong_provider", available_values=["openai"]
)
with patch("goose.utils._create_exchange.build_exchange", side_effect=expected_error):
create_exchange(profile=TEST_PROFILE, notifier=TEST_NOTIFIER)
assert "Unknown provider: wrong_provider. Available providers: openai" in mock_print.call_args_list[0][0][0]
mock_sys_exit.assert_called_once_with(1)
class TestWhenProviderEnvVarNotFound:
API_KEY_ENV_VAR = "OPENAI_API_KEY"
API_KEY_ENV_VALUE = "api_key_value"
PROVIDER_NAME = "openai"
SERVICE_NAME = "goose"
EXPECTED_ERROR = MissingProviderEnvVariableError(env_variable=API_KEY_ENV_VAR, provider=PROVIDER_NAME)
def test_create_exchange_get_api_key_from_keychain(
self, mock_print, mock_sys_exit, mock_keyring_get_password, mock_keyring_set_password
):
self._clean_env()
with patch("goose.utils._create_exchange.build_exchange", side_effect=[self.EXPECTED_ERROR, TEST_EXCHANGE]):
mock_keyring_get_password.return_value = self.API_KEY_ENV_VALUE
assert create_exchange(profile=TEST_PROFILE, notifier=TEST_NOTIFIER) == TEST_EXCHANGE
assert os.environ[self.API_KEY_ENV_VAR] == self.API_KEY_ENV_VALUE
mock_keyring_get_password.assert_called_once_with(self.SERVICE_NAME, self.API_KEY_ENV_VAR)
mock_print.assert_called_once_with(
f"Using {self.API_KEY_ENV_VAR} value for {self.PROVIDER_NAME} from your keychain"
)
mock_sys_exit.assert_not_called()
mock_keyring_set_password.assert_not_called()
def test_create_exchange_ask_api_key_and_user_set_in_keychain(
self, mock_prompt, mock_confirm, mock_sys_exit, mock_keyring_get_password, mock_keyring_set_password, mock_print
):
self._clean_env()
with patch("goose.utils._create_exchange.build_exchange", side_effect=[self.EXPECTED_ERROR, TEST_EXCHANGE]):
mock_keyring_get_password.return_value = None
mock_prompt.return_value = self.API_KEY_ENV_VALUE
mock_confirm.return_value = True
assert create_exchange(profile=TEST_NOTIFIER, notifier=TEST_NOTIFIER) == TEST_EXCHANGE
assert os.environ[self.API_KEY_ENV_VAR] == self.API_KEY_ENV_VALUE
mock_keyring_set_password.assert_called_once_with(
self.SERVICE_NAME, self.API_KEY_ENV_VAR, self.API_KEY_ENV_VALUE
)
mock_confirm.assert_called_once_with(
f"Would you like to save the {self.API_KEY_ENV_VAR} value to your keychain?"
)
mock_print.assert_called_once_with(
f"Saved {self.API_KEY_ENV_VAR} to your key_chain. "
+ f"service_name: goose, user_name: {self.API_KEY_ENV_VAR}"
)
mock_sys_exit.assert_not_called()
def test_create_exchange_ask_api_key_and_user_not_set_in_keychain(
self, mock_prompt, mock_confirm, mock_sys_exit, mock_keyring_get_password, mock_keyring_set_password
):
self._clean_env()
with patch("goose.utils._create_exchange.build_exchange", side_effect=[self.EXPECTED_ERROR, TEST_EXCHANGE]):
mock_keyring_get_password.return_value = None
mock_prompt.return_value = self.API_KEY_ENV_VALUE
mock_confirm.return_value = False
assert create_exchange(profile=TEST_NOTIFIER, notifier=TEST_NOTIFIER) == TEST_EXCHANGE
assert os.environ[self.API_KEY_ENV_VAR] == self.API_KEY_ENV_VALUE
mock_keyring_set_password.assert_not_called()
mock_sys_exit.assert_not_called()
def test_create_exchange_fails_when_user_not_provide_api_key(
self, mock_prompt, mock_confirm, mock_sys_exit, mock_keyring_get_password, mock_print
):
self._clean_env()
with patch("goose.utils._create_exchange.build_exchange", side_effect=self.EXPECTED_ERROR):
mock_keyring_get_password.return_value = None
mock_prompt.return_value = None
mock_confirm.return_value = False
create_exchange(profile=TEST_NOTIFIER, notifier=TEST_NOTIFIER)
assert (
"Please set the required environment variable to continue."
in mock_print.call_args_list[0][0][0].renderable
)
mock_sys_exit.assert_called_once_with(1)
def _clean_env(self):
os.environ.pop(self.API_KEY_ENV_VAR, None)