Merge pull request #93 from victorskl/auto-login-no-cache

Auto login when no cache files
victorskl · Mar 3, 2024 · 5a838db · 5a838db
2 parents caabaa1 + 9cacd51
commit 5a838db
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 6 deletions.
diff --git a/tests/test_cmd.py b/tests/test_cmd.py
@@ -39,10 +39,20 @@ def test_is_sso_cached_login_expired_none(self):
         """
         python -m unittest tests.test_cmd.AutoCommandUnitTests.test_is_sso_cached_login_expired_none
         """
-        when(cli.core).get_aws_cli_v2_sso_cached_login(...).thenReturn(None)
-        with ArgvContext(program, '-t', 'auto', '--profile', 'dev'), self.assertRaises(SystemExit) as x:
+        mock_poll = mock(cli.utils.Poll)
+        when(cli.utils).Poll(contains('aws sso login'), ...).thenReturn(mock_poll)
+        when(mock_poll).start(...).thenReturn(mock_poll)
+        when(mock_poll).resolve(...).thenReturn(True)
+
+        when(cli.cmd.AutoCommand).get_aws_cli_v2_sso_cached_login(...).thenReturn(None)
+
+        with ArgvContext(program, '-t', 'auto', '--profile', 'dev'):
             cli.main()
-        self.assertEqual(x.exception.code, 1)
+        cred = cli.utils.read_config(self.credentials.name)
+        new_tok = cred['dev']['aws_session_token']
+        self.assertNotEqual(new_tok, 'tok')
+        self.assertEqual(new_tok, 'VeryLongBase664String==')
+        verify(cli.utils, times=1).Poll(...)
 
     def test_auto_command(self):
         """

diff --git a/yawsso/cmd.py b/yawsso/cmd.py
@@ -138,16 +138,24 @@ def session_cached(self, profile, cached_login):
     def session_refresh(self, profile, cached_login):
         return core.session_refresh(self.login_profile, profile, cached_login)
 
+    @staticmethod
+    def get_aws_cli_v2_sso_cached_login(profile):
+        return core.get_aws_cli_v2_sso_cached_login(profile)
+
     def perform(self):
         profile = core.load_profile_from_config(self.login_profile, self.co.config)
 
         if not core.is_sso_profile(profile):
             utils.halt(f"Login profile is not an AWS SSO profile. Abort auto syncing profile `{self.login_profile}`")
 
-        cached_login = core.get_aws_cli_v2_sso_cached_login(profile)
+        cached_login = self.get_aws_cli_v2_sso_cached_login(profile)
         if cached_login is None:
-            utils.halt(f"Can not find SSO login session cache in {core.aws_sso_cache_path} "
-                       f"for ({profile['sso_start_url']}) profile `{self.login_profile})`.")
+            # The sso cache files may be cleared away by user or some process such as brew. Since we have proper
+            # SSO login profile, just try to auto login. See https://github.com/victorskl/yawsso/issues/87
+            logger.log(TRACE, f"Can not find SSO login session cache in {core.aws_sso_cache_path} "
+                              f"for {profile['sso_start_url']} profile `{self.login_profile}`.")
+            super(AutoCommand, self).perform()
+            return self
 
         # try 1: attempt using cached accessToken
         role_cred_success, _ = self.session_cached(profile, cached_login)