Skip to content

Commit

Permalink
openai[patch]: add tests for secret_str for keys (#20982)
Browse files Browse the repository at this point in the history
**Description:** Add tests to check API keys and Active Directory tokens
are masked
**Issue:** Resolves #12165 for OpenAI and Azure OpenAI models
**Dependencies:** None

Also resolves #12473 which may be closed.

Additional contributors @alex4321 (#12473) and @onesolpark (#12542)
  • Loading branch information
sepiatone committed May 1, 2024
1 parent 45ddf4d commit 12b1caf
Showing 1 changed file with 127 additions and 0 deletions.
127 changes: 127 additions & 0 deletions libs/partners/openai/tests/unit_tests/test_secrets.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,9 @@
from typing import Type, cast

import pytest
from langchain_core.pydantic_v1 import SecretStr
from pytest import CaptureFixture, MonkeyPatch

from langchain_openai import (
AzureChatOpenAI,
AzureOpenAI,
Expand Down Expand Up @@ -60,3 +66,124 @@ def test_azure_openai_embeddings_secrets() -> None:
s = str(o)
assert "foo1" not in s
assert "foo2" not in s


@pytest.mark.parametrize(
"model_class", [AzureChatOpenAI, AzureOpenAI, AzureOpenAIEmbeddings]
)
def test_azure_openai_api_key_is_secret_string(model_class: Type) -> None:
"""Test that the API key is stored as a SecretStr."""
model = model_class(
openai_api_key="secret-api-key",
azure_endpoint="endpoint",
azure_ad_token="secret-ad-token",
api_version="version",
)
assert isinstance(model.openai_api_key, SecretStr)
assert isinstance(model.azure_ad_token, SecretStr)


@pytest.mark.parametrize(
"model_class", [AzureChatOpenAI, AzureOpenAI, AzureOpenAIEmbeddings]
)
def test_azure_openai_api_key_masked_when_passed_from_env(
model_class: Type, monkeypatch: MonkeyPatch, capsys: CaptureFixture
) -> None:
"""Test that the API key is masked when passed from an environment variable."""
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "secret-api-key")
monkeypatch.setenv("AZURE_OPENAI_AD_TOKEN", "secret-ad-token")
model = model_class(
azure_endpoint="endpoint",
api_version="version",
)
print(model.openai_api_key, end="") # noqa: T201
captured = capsys.readouterr()

assert captured.out == "**********"

print(model.azure_ad_token, end="") # noqa: T201
captured = capsys.readouterr()

assert captured.out == "**********"


@pytest.mark.parametrize(
"model_class", [AzureChatOpenAI, AzureOpenAI, AzureOpenAIEmbeddings]
)
def test_azure_openai_api_key_masked_when_passed_via_constructor(
model_class: Type,
capsys: CaptureFixture,
) -> None:
"""Test that the API key is masked when passed via the constructor."""
model = model_class(
openai_api_key="secret-api-key",
azure_endpoint="endpoint",
azure_ad_token="secret-ad-token",
api_version="version",
)
print(model.openai_api_key, end="") # noqa: T201
captured = capsys.readouterr()

assert captured.out == "**********"

print(model.azure_ad_token, end="") # noqa: T201
captured = capsys.readouterr()

assert captured.out == "**********"


@pytest.mark.parametrize(
"model_class", [AzureChatOpenAI, AzureOpenAI, AzureOpenAIEmbeddings]
)
def test_azure_openai_uses_actual_secret_value_from_secretstr(
model_class: Type,
) -> None:
"""Test that the actual secret value is correctly retrieved."""
model = model_class(
openai_api_key="secret-api-key",
azure_endpoint="endpoint",
azure_ad_token="secret-ad-token",
api_version="version",
)
assert cast(SecretStr, model.openai_api_key).get_secret_value() == "secret-api-key"
assert cast(SecretStr, model.azure_ad_token).get_secret_value() == "secret-ad-token"


@pytest.mark.parametrize("model_class", [ChatOpenAI, OpenAI, OpenAIEmbeddings])
def test_openai_api_key_is_secret_string(model_class: Type) -> None:
"""Test that the API key is stored as a SecretStr."""
model = model_class(openai_api_key="secret-api-key")
assert isinstance(model.openai_api_key, SecretStr)


@pytest.mark.parametrize("model_class", [ChatOpenAI, OpenAI, OpenAIEmbeddings])
def test_openai_api_key_masked_when_passed_from_env(
model_class: Type, monkeypatch: MonkeyPatch, capsys: CaptureFixture
) -> None:
"""Test that the API key is masked when passed from an environment variable."""
monkeypatch.setenv("OPENAI_API_KEY", "secret-api-key")
model = model_class()
print(model.openai_api_key, end="") # noqa: T201
captured = capsys.readouterr()

assert captured.out == "**********"


@pytest.mark.parametrize("model_class", [ChatOpenAI, OpenAI, OpenAIEmbeddings])
def test_openai_api_key_masked_when_passed_via_constructor(
model_class: Type,
capsys: CaptureFixture,
) -> None:
"""Test that the API key is masked when passed via the constructor."""
model = model_class(openai_api_key="secret-api-key")
print(model.openai_api_key, end="") # noqa: T201
captured = capsys.readouterr()

assert captured.out == "**********"


@pytest.mark.parametrize("model_class", [ChatOpenAI, OpenAI, OpenAIEmbeddings])
def test_openai_uses_actual_secret_value_from_secretstr(model_class: Type) -> None:
"""Test that the actual secret value is correctly retrieved."""
model = model_class(openai_api_key="secret-api-key")
assert cast(SecretStr, model.openai_api_key).get_secret_value() == "secret-api-key"

0 comments on commit 12b1caf

Please sign in to comment.