diff --git a/lib/crewai-tools/tests/test_platform_tools_misc.py b/lib/crewai-tools/tests/test_platform_tools_misc.py index d36c1ded2..c3101cd2f 100644 --- a/lib/crewai-tools/tests/test_platform_tools_misc.py +++ b/lib/crewai-tools/tests/test_platform_tools_misc.py @@ -4,7 +4,7 @@ import os from unittest.mock import patch import pytest -from crewai.context import platform_integration_context +from crewai.context import platform_integration_context, set_platform_integration_token, reset_platform_integration_token from crewai_tools.tools.crewai_platform_tools.misc import ( get_platform_integration_token, ) @@ -14,6 +14,17 @@ from crewai_tools.tools.crewai_platform_tools.misc import ( class TestTokenRetrievalWithFallback: """Test token retrieval logic with environment fallback.""" + @pytest.fixture + def clean_context(self): + token = set_platform_integration_token(None) + env_backup = os.environ.pop("CREWAI_PLATFORM_INTEGRATION_TOKEN", None) + yield + reset_platform_integration_token(token) + if env_backup is not None: + os.environ["CREWAI_PLATFORM_INTEGRATION_TOKEN"] = env_backup + else: + os.environ.pop("CREWAI_PLATFORM_INTEGRATION_TOKEN", None) + def test_context_token_takes_precedence(self, clean_context): """Test that context token takes precedence over environment variable.""" context_token = "context-token" diff --git a/lib/crewai/src/crewai/context.py b/lib/crewai/src/crewai/context.py index 149a85d42..4d6338ce2 100644 --- a/lib/crewai/src/crewai/context.py +++ b/lib/crewai/src/crewai/context.py @@ -1,7 +1,7 @@ from collections.abc import Generator -from contextlib import contextmanager, nullcontext +from contextlib import AbstractContextManager, contextmanager, nullcontext import contextvars -from typing import Any, ContextManager +from typing import Any _platform_integration_token: contextvars.ContextVar[str | None] = ( @@ -31,7 +31,7 @@ def get_platform_integration_token() -> str | None: return _platform_integration_token.get() -def platform_integration_context(integration_token: str | None) -> ContextManager[None]: +def platform_integration_context(integration_token: str | None) -> AbstractContextManager[None]: """Context manager to temporarily set the platform integration token. Args: