Make to preserve config value if exists

Author: code-yeongyu

aishell/cli/config_aishell.py

@@ -11,17 +11,19 @@
 
 
 def config_aishell():
-    rich.print('''Hi! 🙌 I am [bold blue]AiShell[/bold blue], [yellow]your powerful terminal assistant[/yellow] 🔥
+    rich.print('''
+Hi! 🙌 I am [bold blue]AiShell[/bold blue], [yellow]your powerful terminal assistant[/yellow] 🔥
 I am here to assist you with configuring AiShell. 💪
 
 Please make sure that you have logged into chat.openai.com on your browser before we continue. 🗝️
 
-''')
+'''[1:])
     typer.confirm('Are you ready to proceed? 🚀', abort=True)
 
-    rich.print(f'''Which browser did you use to log in to chat.openai.com?
+    rich.print(f'''
+Which browser did you use to log in to chat.openai.com?
 
-We support the following browsers: [{SUPPORTED_BROWSERS}]''')
+We support the following browsers: [{SUPPORTED_BROWSERS}]'''[1:])
     browser_name = typer.prompt('Please enter your choice here: ')
     if browser_name not in SUPPORTED_BROWSERS:
         rich.print(f'Browser {browser_name} is not supported. Supported browsers are: {SUPPORTED_BROWSERS}')
@@ -33,18 +35,33 @@ def config_aishell():
         rich.print('Failed to get session token. 😓 Can you check if you are logged in to https://chat.openai.com?')
         sys.exit(1)
 
-    is_paid = typer.confirm("It's my last question! 🤩 Are you a PLUS user?")
+    config_manager = save_config(session_token)
 
-    chatgpt_config = RevChatGPTChatbotConfigModel(session_token=session_token, paid=is_paid)
-    aishell_config = AiShellConfigModel(chatgpt_config=chatgpt_config)
-    config_manager = AiShellConfigManager(config_model=aishell_config)
-    config_manager.save_config()
-
-    rich.print(f'''[green bold]Excellent![/green bold] You are now ready to use [bold blue]AiShell[/bold blue] 🚀
+    rich.print(f'''
+[green bold]Excellent![/green bold] You are now ready to use [bold blue]AiShell[/bold blue] 🚀
 
 Enjoy your AI powered terminal assistant! 🎉
 
 [dim]To check your settings file, it's at: {config_manager.config_path}[/dim]
 
-''')
+'''[1:])
+    return config_manager
+
+
+def save_config(session_token: str):
+    is_config_file_available = AiShellConfigManager.is_config_file_available(AiShellConfigManager.DEFAULT_CONFIG_PATH)
+    if is_config_file_available:
+        config_manager = AiShellConfigManager(load_config=True)
+        is_chatgpt_config_available = config_manager.config_model.chatgpt_config is not None
+        if is_chatgpt_config_available:
+            assert config_manager.config_model.chatgpt_config  # for type hinting
+            config_manager.config_model.chatgpt_config.session_token = session_token
+        else:
+            config_manager.config_model.chatgpt_config = RevChatGPTChatbotConfigModel(session_token=session_token)
+    else:
+        chatgpt_config = RevChatGPTChatbotConfigModel(session_token=session_token)
+        aishell_config = AiShellConfigModel(chatgpt_config=chatgpt_config)
+        config_manager = AiShellConfigManager(config_model=aishell_config)
+
+    config_manager.save_config()
     return config_manager

aishell/cli/test/__init__.py

@@ -0,0 +1,63 @@
+from unittest.mock import MagicMock, patch
+
+from aishell.cli.config_aishell import save_config
+from aishell.models import AiShellConfigModel, LanguageModel, RevChatGPTChatbotConfigModel
+from aishell.utils import AiShellConfigManager
+
+
+class TestSaveConfig:
+
+    @patch('aishell.cli.config_aishell.AiShellConfigManager', spec=AiShellConfigManager)
+    def test_success_when_config_file_not_available(
+        self,
+        mocked_config_manager_class: MagicMock,
+    ):
+        '''config file 이 없는 경우 테스트 성공해야 한다'''
+        # given
+        mocked_config_manager_class.is_config_file_available.return_value = False
+
+        # when
+        save_config('valid_session_token')
+
+        # then
+        mocked_config_manager_class.return_value.save_config.assert_called_once()
+
+    @patch('aishell.cli.config_aishell.AiShellConfigManager', spec=AiShellConfigManager)
+    def test_success_when_config_file_available_with_chatgpt_config(
+        self,
+        mocked_config_manager_class: MagicMock,
+    ):
+        '''config file 이 있고, chatgpt_config 가 있는 경우 테스트 성공해야 한다.'''
+        # given
+        mocked_config_manager_class.is_config_file_available.return_value = True
+        mocked_config_manager_instance = mocked_config_manager_class.return_value
+        mocked_config_manager_instance.config_model =\
+            AiShellConfigModel(chatgpt_config=RevChatGPTChatbotConfigModel(session_token='invalid_session_token'))
+
+        # when
+        save_config('valid_session_token')
+
+        # then
+        mocked_config_manager_class.return_value.save_config.assert_called_once()
+
+    @patch('aishell.cli.config_aishell.AiShellConfigManager', spec=AiShellConfigManager)
+    def test_success_when_config_file_available_without_chatgpt_config(
+        self,
+        mocked_config_manager_class: MagicMock,
+    ):
+        '''config file 이 있고, chatgpt_config 가 없는 경우 테스트 성공해야 한다.'''
+        # given
+        mocked_config_manager_class.is_config_file_available.return_value = True
+        mocked_config_manager_instance = mocked_config_manager_class.return_value
+        mocked_config_manager_instance.config_model =\
+            AiShellConfigModel(language_model=LanguageModel.GPT3, openai_api_key='valid_openai_api_key')
+
+        # when
+        save_config('valid_session_token')
+
+        # then
+        mocked_config_manager_class.return_value.save_config.assert_called_once()
+
+    # when config file available - with chatgpt_config
+    # when config file available - without chatgpt_config
+    # when config file not available