diff --git a/tests/unit_tests/test_llm_providers.py b/tests/unit_tests/test_llm_providers.py index 93ab66cff..43525c23a 100644 --- a/tests/unit_tests/test_llm_providers.py +++ b/tests/unit_tests/test_llm_providers.py @@ -1,4 +1,5 @@ import importlib.util +import os from typing import Any, Callable, Dict, Iterable from unittest.mock import MagicMock @@ -14,6 +15,7 @@ get_llm_ask, ) from guardrails.utils.openai_utils import OPENAI_VERSION +from guardrails.utils.safe_get import safe_get_with_brackets from .mocks import MockAsyncOpenAILlm, MockOpenAILlm @@ -593,7 +595,14 @@ def test_get_llm_ask_openai_chat(): not importlib.util.find_spec("manifest"), reason="manifest is not installed", ) -def test_get_llm_ask_manifest(): +def test_get_llm_ask_manifest(mocker): + def mock_os_environ_get(key, *args): + if key == "OPENAI_API_KEY": + return "sk-xxxxxxxxxxxxxx" + return safe_get_with_brackets(os.environ, key, *args) + + mocker.patch("os.environ.get", side_effect=mock_os_environ_get) + from manifest import Manifest from guardrails.llm_providers import ManifestCallable