feat: reset tokens on crewai config reset (#3365)

tests/cli/authentication/test_utils.py

@@ -1,17 +1,14 @@
-import json
 import jwt
 import unittest
-from datetime import datetime, timedelta
 from unittest.mock import MagicMock, patch
 
-from cryptography.fernet import Fernet
 
-from crewai.cli.authentication.utils import TokenManager, validate_jwt_token
+from crewai.cli.authentication.utils import validate_jwt_token
 
 
 @patch("crewai.cli.authentication.utils.PyJWKClient", return_value=MagicMock())
 @patch("crewai.cli.authentication.utils.jwt")
-class TestValidateToken(unittest.TestCase):
+class TestUtils(unittest.TestCase):
     def test_validate_jwt_token(self, mock_jwt, mock_pyjwkclient):
         mock_jwt.decode.return_value = {"exp": 1719859200}
 
@@ -105,121 +102,3 @@ class TestValidateToken(unittest.TestCase):
                 issuer="https://mock_issuer",
                 audience="app_id_xxxx",
             )
-
-
-class TestTokenManager(unittest.TestCase):
-    @patch("crewai.cli.authentication.utils.TokenManager._get_or_create_key")
-    def setUp(self, mock_get_key):
-        mock_get_key.return_value = Fernet.generate_key()
-        self.token_manager = TokenManager()
-
-    @patch("crewai.cli.authentication.utils.TokenManager.read_secure_file")
-    @patch("crewai.cli.authentication.utils.TokenManager.save_secure_file")
-    @patch("crewai.cli.authentication.utils.TokenManager._get_or_create_key")
-    def test_get_or_create_key_existing(self, mock_get_or_create, mock_save, mock_read):
-        mock_key = Fernet.generate_key()
-        mock_get_or_create.return_value = mock_key
-
-        token_manager = TokenManager()
-        result = token_manager.key
-
-        self.assertEqual(result, mock_key)
-
-    @patch("crewai.cli.authentication.utils.Fernet.generate_key")
-    @patch("crewai.cli.authentication.utils.TokenManager.read_secure_file")
-    @patch("crewai.cli.authentication.utils.TokenManager.save_secure_file")
-    def test_get_or_create_key_new(self, mock_save, mock_read, mock_generate):
-        mock_key = b"new_key"
-        mock_read.return_value = None
-        mock_generate.return_value = mock_key
-
-        result = self.token_manager._get_or_create_key()
-
-        self.assertEqual(result, mock_key)
-        mock_read.assert_called_once_with("secret.key")
-        mock_generate.assert_called_once()
-        mock_save.assert_called_once_with("secret.key", mock_key)
-
-    @patch("crewai.cli.authentication.utils.TokenManager.save_secure_file")
-    def test_save_tokens(self, mock_save):
-        access_token = "test_token"
-        expires_at = int((datetime.now() + timedelta(seconds=3600)).timestamp())
-
-        self.token_manager.save_tokens(access_token, expires_at)
-
-        mock_save.assert_called_once()
-        args = mock_save.call_args[0]
-        self.assertEqual(args[0], "tokens.enc")
-        decrypted_data = self.token_manager.fernet.decrypt(args[1])
-        data = json.loads(decrypted_data)
-        self.assertEqual(data["access_token"], access_token)
-        expiration = datetime.fromisoformat(data["expiration"])
-        self.assertEqual(expiration, datetime.fromtimestamp(expires_at))
-
-    @patch("crewai.cli.authentication.utils.TokenManager.read_secure_file")
-    def test_get_token_valid(self, mock_read):
-        access_token = "test_token"
-        expiration = (datetime.now() + timedelta(hours=1)).isoformat()
-        data = {"access_token": access_token, "expiration": expiration}
-        encrypted_data = self.token_manager.fernet.encrypt(json.dumps(data).encode())
-        mock_read.return_value = encrypted_data
-
-        result = self.token_manager.get_token()
-
-        self.assertEqual(result, access_token)
-
-    @patch("crewai.cli.authentication.utils.TokenManager.read_secure_file")
-    def test_get_token_expired(self, mock_read):
-        access_token = "test_token"
-        expiration = (datetime.now() - timedelta(hours=1)).isoformat()
-        data = {"access_token": access_token, "expiration": expiration}
-        encrypted_data = self.token_manager.fernet.encrypt(json.dumps(data).encode())
-        mock_read.return_value = encrypted_data
-
-        result = self.token_manager.get_token()
-
-        self.assertIsNone(result)
-
-    @patch("crewai.cli.authentication.utils.TokenManager.get_secure_storage_path")
-    @patch("builtins.open", new_callable=unittest.mock.mock_open)
-    @patch("crewai.cli.authentication.utils.os.chmod")
-    def test_save_secure_file(self, mock_chmod, mock_open, mock_get_path):
-        mock_path = MagicMock()
-        mock_get_path.return_value = mock_path
-        filename = "test_file.txt"
-        content = b"test_content"
-
-        self.token_manager.save_secure_file(filename, content)
-
-        mock_path.__truediv__.assert_called_once_with(filename)
-        mock_open.assert_called_once_with(mock_path.__truediv__.return_value, "wb")
-        mock_open().write.assert_called_once_with(content)
-        mock_chmod.assert_called_once_with(mock_path.__truediv__.return_value, 0o600)
-
-    @patch("crewai.cli.authentication.utils.TokenManager.get_secure_storage_path")
-    @patch(
-        "builtins.open", new_callable=unittest.mock.mock_open, read_data=b"test_content"
-    )
-    def test_read_secure_file_exists(self, mock_open, mock_get_path):
-        mock_path = MagicMock()
-        mock_get_path.return_value = mock_path
-        mock_path.__truediv__.return_value.exists.return_value = True
-        filename = "test_file.txt"
-
-        result = self.token_manager.read_secure_file(filename)
-
-        self.assertEqual(result, b"test_content")
-        mock_path.__truediv__.assert_called_once_with(filename)
-        mock_open.assert_called_once_with(mock_path.__truediv__.return_value, "rb")
-
-    @patch("crewai.cli.authentication.utils.TokenManager.get_secure_storage_path")
-    def test_read_secure_file_not_exists(self, mock_get_path):
-        mock_path = MagicMock()
-        mock_get_path.return_value = mock_path
-        mock_path.__truediv__.return_value.exists.return_value = False
-        filename = "test_file.txt"
-
-        result = self.token_manager.read_secure_file(filename)
-
-        self.assertIsNone(result)
-        mock_path.__truediv__.assert_called_once_with(filename)

tests/cli/test_config.py

@@ -3,6 +3,7 @@ import shutil
 import tempfile
 import unittest
 from pathlib import Path
+from unittest.mock import patch, MagicMock
 
 from crewai.cli.config import (
     Settings,
@@ -10,6 +11,8 @@ from crewai.cli.config import (
     CLI_SETTINGS_KEYS,
     DEFAULT_CLI_SETTINGS,
 )
+from crewai.cli.shared.token_manager import TokenManager
+from datetime import datetime, timedelta
 
 
 class TestSettings(unittest.TestCase):
@@ -66,7 +69,8 @@ class TestSettings(unittest.TestCase):
         for key in user_settings.keys():
             self.assertEqual(getattr(settings, key), None)
 
-    def test_reset_settings(self):
+    @patch("crewai.cli.config.TokenManager")
+    def test_reset_settings(self, mock_token_manager):
         user_settings = {key: f"value_for_{key}" for key in USER_SETTINGS_KEYS}
         cli_settings = {key: f"value_for_{key}" for key in CLI_SETTINGS_KEYS}
 
@@ -74,6 +78,11 @@ class TestSettings(unittest.TestCase):
             config_path=self.config_path, **user_settings, **cli_settings
         )
 
+        mock_token_manager.return_value = MagicMock()
+        TokenManager().save_tokens(
+            "aaa.bbb.ccc", (datetime.now() + timedelta(seconds=36000)).timestamp()
+        )
+
         settings.reset()
 
         for key in user_settings.keys():
@@ -81,6 +90,8 @@ class TestSettings(unittest.TestCase):
         for key in cli_settings.keys():
             self.assertEqual(getattr(settings, key), DEFAULT_CLI_SETTINGS.get(key))
 
+        mock_token_manager.return_value.clear_tokens.assert_called_once()
+
     def test_dump_new_settings(self):
         settings = Settings(
             config_path=self.config_path, tool_repository_username="user1"

tests/cli/test_token_manager.py

@@ -0,0 +1,138 @@
+import json
+import unittest
+from datetime import datetime, timedelta
+from unittest.mock import MagicMock, patch
+
+from cryptography.fernet import Fernet
+
+from crewai.cli.shared.token_manager import TokenManager
+
+
+class TestTokenManager(unittest.TestCase):
+    @patch("crewai.cli.shared.token_manager.TokenManager._get_or_create_key")
+    def setUp(self, mock_get_key):
+        mock_get_key.return_value = Fernet.generate_key()
+        self.token_manager = TokenManager()
+
+    @patch("crewai.cli.shared.token_manager.TokenManager.read_secure_file")
+    @patch("crewai.cli.shared.token_manager.TokenManager.save_secure_file")
+    @patch("crewai.cli.shared.token_manager.TokenManager._get_or_create_key")
+    def test_get_or_create_key_existing(self, mock_get_or_create, mock_save, mock_read):
+        mock_key = Fernet.generate_key()
+        mock_get_or_create.return_value = mock_key
+
+        token_manager = TokenManager()
+        result = token_manager.key
+
+        self.assertEqual(result, mock_key)
+
+    @patch("crewai.cli.shared.token_manager.Fernet.generate_key")
+    @patch("crewai.cli.shared.token_manager.TokenManager.read_secure_file")
+    @patch("crewai.cli.shared.token_manager.TokenManager.save_secure_file")
+    def test_get_or_create_key_new(self, mock_save, mock_read, mock_generate):
+        mock_key = b"new_key"
+        mock_read.return_value = None
+        mock_generate.return_value = mock_key
+
+        result = self.token_manager._get_or_create_key()
+
+        self.assertEqual(result, mock_key)
+        mock_read.assert_called_once_with("secret.key")
+        mock_generate.assert_called_once()
+        mock_save.assert_called_once_with("secret.key", mock_key)
+
+    @patch("crewai.cli.shared.token_manager.TokenManager.save_secure_file")
+    def test_save_tokens(self, mock_save):
+        access_token = "test_token"
+        expires_at = int((datetime.now() + timedelta(seconds=3600)).timestamp())
+
+        self.token_manager.save_tokens(access_token, expires_at)
+
+        mock_save.assert_called_once()
+        args = mock_save.call_args[0]
+        self.assertEqual(args[0], "tokens.enc")
+        decrypted_data = self.token_manager.fernet.decrypt(args[1])
+        data = json.loads(decrypted_data)
+        self.assertEqual(data["access_token"], access_token)
+        expiration = datetime.fromisoformat(data["expiration"])
+        self.assertEqual(expiration, datetime.fromtimestamp(expires_at))
+
+    @patch("crewai.cli.shared.token_manager.TokenManager.read_secure_file")
+    def test_get_token_valid(self, mock_read):
+        access_token = "test_token"
+        expiration = (datetime.now() + timedelta(hours=1)).isoformat()
+        data = {"access_token": access_token, "expiration": expiration}
+        encrypted_data = self.token_manager.fernet.encrypt(json.dumps(data).encode())
+        mock_read.return_value = encrypted_data
+
+        result = self.token_manager.get_token()
+
+        self.assertEqual(result, access_token)
+
+    @patch("crewai.cli.shared.token_manager.TokenManager.read_secure_file")
+    def test_get_token_expired(self, mock_read):
+        access_token = "test_token"
+        expiration = (datetime.now() - timedelta(hours=1)).isoformat()
+        data = {"access_token": access_token, "expiration": expiration}
+        encrypted_data = self.token_manager.fernet.encrypt(json.dumps(data).encode())
+        mock_read.return_value = encrypted_data
+
+        result = self.token_manager.get_token()
+
+        self.assertIsNone(result)
+
+    @patch("crewai.cli.shared.token_manager.TokenManager.get_secure_storage_path")
+    @patch("builtins.open", new_callable=unittest.mock.mock_open)
+    @patch("crewai.cli.shared.token_manager.os.chmod")
+    def test_save_secure_file(self, mock_chmod, mock_open, mock_get_path):
+        mock_path = MagicMock()
+        mock_get_path.return_value = mock_path
+        filename = "test_file.txt"
+        content = b"test_content"
+
+        self.token_manager.save_secure_file(filename, content)
+
+        mock_path.__truediv__.assert_called_once_with(filename)
+        mock_open.assert_called_once_with(mock_path.__truediv__.return_value, "wb")
+        mock_open().write.assert_called_once_with(content)
+        mock_chmod.assert_called_once_with(mock_path.__truediv__.return_value, 0o600)
+
+    @patch("crewai.cli.shared.token_manager.TokenManager.get_secure_storage_path")
+    @patch(
+        "builtins.open", new_callable=unittest.mock.mock_open, read_data=b"test_content"
+    )
+    def test_read_secure_file_exists(self, mock_open, mock_get_path):
+        mock_path = MagicMock()
+        mock_get_path.return_value = mock_path
+        mock_path.__truediv__.return_value.exists.return_value = True
+        filename = "test_file.txt"
+
+        result = self.token_manager.read_secure_file(filename)
+
+        self.assertEqual(result, b"test_content")
+        mock_path.__truediv__.assert_called_once_with(filename)
+        mock_open.assert_called_once_with(mock_path.__truediv__.return_value, "rb")
+
+    @patch("crewai.cli.shared.token_manager.TokenManager.get_secure_storage_path")
+    def test_read_secure_file_not_exists(self, mock_get_path):
+        mock_path = MagicMock()
+        mock_get_path.return_value = mock_path
+        mock_path.__truediv__.return_value.exists.return_value = False
+        filename = "test_file.txt"
+
+        result = self.token_manager.read_secure_file(filename)
+
+        self.assertIsNone(result)
+        mock_path.__truediv__.assert_called_once_with(filename)
+
+    @patch("crewai.cli.shared.token_manager.TokenManager.get_secure_storage_path")
+    def test_clear_tokens(self, mock_get_path):
+        mock_path = MagicMock()
+        mock_get_path.return_value = mock_path
+
+        self.token_manager.clear_tokens()
+
+        mock_path.__truediv__.assert_called_once_with("tokens.enc")
+        mock_path.__truediv__.return_value.unlink.assert_called_once_with(
+            missing_ok=True
+        )

tests/cli/tools/test_main.py

@@ -11,7 +11,7 @@ from unittest.mock import MagicMock, patch
 import pytest
 from pytest import raises
 
-from crewai.cli.authentication.utils import TokenManager
+from crewai.cli.shared.token_manager import TokenManager
 from crewai.cli.tools.main import ToolCommand