From 12b1caf2959ebe773f91b0e264ae7adfff55f43d Mon Sep 17 00:00:00 2001
From: aditya thomas <aditya.thomas@alum.mit.edu>
Date: Wed, 1 May 2024 10:56:20 +0530
Subject: [PATCH] openai[patch]: add tests for secret_str for keys (#20982)

**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)
---
 .../openai/tests/unit_tests/test_secrets.py   | 127 ++++++++++++++++++
 1 file changed, 127 insertions(+)

diff --git a/libs/partners/openai/tests/unit_tests/test_secrets.py b/libs/partners/openai/tests/unit_tests/test_secrets.py
index a90fbb4b7dbab..928ee5ba735e3 100644
--- a/libs/partners/openai/tests/unit_tests/test_secrets.py
+++ b/libs/partners/openai/tests/unit_tests/test_secrets.py
@@ -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,
@@ -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"