update from review

langchain-ai · Apr 30, 2024 · c498422 · c498422
1 parent 039cefd
commit c498422
Showing 1 changed file with 41 additions and 243 deletions.
diff --git a/libs/partners/openai/tests/unit_tests/test_secrets.py b/libs/partners/openai/tests/unit_tests/test_secrets.py
@@ -1,5 +1,6 @@
-from typing import cast
+from typing import Type, cast
 
+import pytest
 from langchain_core.pydantic_v1 import SecretStr
 from pytest import CaptureFixture, MonkeyPatch
 
@@ -67,10 +68,12 @@ def test_azure_openai_embeddings_secrets() -> None:
     assert "foo2" not in s
 
 
-# azure chat openai
-def test_azure_chat_openai_api_key_is_secret_string() -> None:
+@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."""
-    chat_model = AzureChatOpenAI(
+    chat_model = model_class(
         openai_api_key="secret-api-key",
         azure_endpoint="endpoint",
         azure_ad_token="secret-ad-token",
@@ -80,16 +83,17 @@ def test_azure_chat_openai_api_key_is_secret_string() -> None:
     assert isinstance(chat_model.azure_ad_token, SecretStr)
 
 
-def test_azure_chat_openai_api_key_masked_when_passed_from_env(
-    monkeypatch: MonkeyPatch, capsys: CaptureFixture
+@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")
-    chat_model = AzureChatOpenAI(
-        openai_api_key="secret-api-key",
+    chat_model = model_class(
         azure_endpoint="endpoint",
-        azure_ad_token="secret-ad-token",
         api_version="version",
     )
     print(chat_model.openai_api_key, end="")  # noqa: T201
@@ -103,11 +107,15 @@ def test_azure_chat_openai_api_key_masked_when_passed_from_env(
     assert captured.out == "**********"
 
 
-def test_azure_chat_openai_api_key_masked_when_passed_via_constructor(
+@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."""
-    chat_model = AzureChatOpenAI(
+    chat_model = model_class(
         openai_api_key="secret-api-key",
         azure_endpoint="endpoint",
         azure_ad_token="secret-ad-token",
@@ -124,9 +132,14 @@ def test_azure_chat_openai_api_key_masked_when_passed_via_constructor(
     assert captured.out == "**********"
 
 
-def test_azure_chat_openai_uses_actual_secret_value_from_secretstr() -> None:
+@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."""
-    chat_model = AzureChatOpenAI(
+    chat_model = model_class(
         openai_api_key="secret-api-key",
         azure_endpoint="endpoint",
         azure_ad_token="secret-ad-token",
@@ -142,259 +155,44 @@ def test_azure_chat_openai_uses_actual_secret_value_from_secretstr() -> None:
     )
 
 
-# azure openai
-def test_azure_openai_api_key_is_secret_string() -> None:
-    """Test that the API key is stored as a SecretStr."""
-    llm = AzureOpenAI(
-        openai_api_key="secret-api-key",
-        azure_endpoint="endpoint",
-        azure_ad_token="secret-ad-token",
-        api_version="version",
-    )
-    assert isinstance(llm.openai_api_key, SecretStr)
-    assert isinstance(llm.azure_ad_token, SecretStr)
-
-
-def test_azure_openai_api_key_masked_when_passed_from_env(
-    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")
-    llm = AzureOpenAI(
-        openai_api_key="secret-api-key",
-        azure_endpoint="endpoint",
-        azure_ad_token="secret-ad-token",
-        api_version="version",
-    )
-    print(llm.openai_api_key, end="")  # noqa: T201
-    captured = capsys.readouterr()
-
-    assert captured.out == "**********"
-
-    print(llm.azure_ad_token, end="")  # noqa: T201
-    captured = capsys.readouterr()
-
-    assert captured.out == "**********"
-
-
-def test_azure_openai_api_key_masked_when_passed_via_constructor(
-    capsys: CaptureFixture,
-) -> None:
-    """Test that the API key is masked when passed via the constructor."""
-    llm = AzureOpenAI(
-        openai_api_key="secret-api-key",
-        azure_endpoint="endpoint",
-        azure_ad_token="secret-ad-token",
-        api_version="version",
-    )
-    print(llm.openai_api_key, end="")  # noqa: T201
-    captured = capsys.readouterr()
-
-    assert captured.out == "**********"
-
-    print(llm.azure_ad_token, end="")  # noqa: T201
-    captured = capsys.readouterr()
-
-    assert captured.out == "**********"
-
-
-def test_azure_openai_uses_actual_secret_value_from_secretstr() -> None:
-    """Test that the actual secret value is correctly retrieved."""
-    llm = AzureOpenAI(
-        openai_api_key="secret-api-key",
-        azure_endpoint="endpoint",
-        azure_ad_token="secret-ad-token",
-        api_version="version",
-    )
-    assert cast(SecretStr, llm.openai_api_key).get_secret_value() == "secret-api-key"
-    assert cast(SecretStr, llm.azure_ad_token).get_secret_value() == "secret-ad-token"
-
-
-# azure openai embeddings
-def test_azure_openai_embeddings_api_key_is_secret_string() -> None:
+@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."""
-    embeddings_model = AzureOpenAIEmbeddings(
-        openai_api_key="secret-api-key",
-        azure_endpoint="endpoint",
-        azure_ad_token="secret-ad-token",
-        api_version="version",
-    )
-    assert isinstance(embeddings_model.openai_api_key, SecretStr)
-    assert isinstance(embeddings_model.azure_ad_token, SecretStr)
-
-
-def test_azure_openai_embeddings_api_key_masked_when_passed_from_env(
-    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")
-    embeddings_model = AzureOpenAIEmbeddings(
-        openai_api_key="secret-api-key",
-        azure_endpoint="endpoint",
-        azure_ad_token="secret-ad-token",
-        api_version="version",
-    )
-    print(embeddings_model.openai_api_key, end="")  # noqa: T201
-    captured = capsys.readouterr()
-
-    assert captured.out == "**********"
-
-    print(embeddings_model.azure_ad_token, end="")  # noqa: T201
-    captured = capsys.readouterr()
-
-    assert captured.out == "**********"
-
-
-def test_azure_openai_embeddings_api_key_masked_when_passed_via_constructor(
-    capsys: CaptureFixture,
-) -> None:
-    """Test that the API key is masked when passed via the constructor."""
-    embeddings_model = AzureOpenAIEmbeddings(
-        openai_api_key="secret-api-key",
-        azure_endpoint="endpoint",
-        azure_ad_token="secret-ad-token",
-        api_version="version",
-    )
-    print(embeddings_model.openai_api_key, end="")  # noqa: T201
-    captured = capsys.readouterr()
-
-    assert captured.out == "**********"
-
-    print(embeddings_model.azure_ad_token, end="")  # noqa: T201
-    captured = capsys.readouterr()
-
-    assert captured.out == "**********"
-
-
-def test_azure_openai_embeddings_uses_actual_secret_value_from_secretstr() -> None:
-    """Test that the actual secret value is correctly retrieved."""
-    embeddings_model = AzureOpenAIEmbeddings(
-        openai_api_key="secret-api-key",
-        azure_endpoint="endpoint",
-        azure_ad_token="secret-ad-token",
-        api_version="version",
-    )
-    assert (
-        cast(SecretStr, embeddings_model.openai_api_key).get_secret_value()
-        == "secret-api-key"
-    )
-    assert (
-        cast(SecretStr, embeddings_model.azure_ad_token).get_secret_value()
-        == "secret-ad-token"
-    )
-
-
-# openai chat
-def test_chat_openai_api_key_is_secret_string() -> None:
-    """Test that the API key is stored as a SecretStr."""
-    chat_model = ChatOpenAI(openai_api_key="secret-api-key")
+    chat_model = model_class(openai_api_key="secret-api-key")
     assert isinstance(chat_model.openai_api_key, SecretStr)
 
 
-def test_chat_openai_api_key_masked_when_passed_from_env(
-    monkeypatch: MonkeyPatch, capsys: CaptureFixture
+@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("AZURE_OPENAI_API_KEY", "secret-api-key")
-    chat_model = ChatOpenAI(openai_api_key="secret-api-key")
+    monkeypatch.setenv("OPENAI_API_KEY", "secret-api-key")
+    chat_model = model_class()
     print(chat_model.openai_api_key, end="")  # noqa: T201
     captured = capsys.readouterr()
 
     assert captured.out == "**********"
 
 
-def test_chat_openai_api_key_masked_when_passed_via_constructor(
+@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."""
-    chat_model = ChatOpenAI(openai_api_key="secret-api-key")
+    chat_model = model_class(openai_api_key="secret-api-key")
     print(chat_model.openai_api_key, end="")  # noqa: T201
     captured = capsys.readouterr()
 
     assert captured.out == "**********"
 
 
-def test_chat_openai_uses_actual_secret_value_from_secretstr() -> None:
+@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."""
-    chat_model = ChatOpenAI(openai_api_key="secret-api-key")
+    chat_model = model_class(openai_api_key="secret-api-key")
     assert (
         cast(SecretStr, chat_model.openai_api_key).get_secret_value()
         == "secret-api-key"
     )
-
-
-# openai
-def test_openai_api_key_is_secret_string() -> None:
-    """Test that the API key is stored as a SecretStr."""
-    llm = OpenAI(openai_api_key="secret-api-key")
-    assert isinstance(llm.openai_api_key, SecretStr)
-
-
-def test_openai_api_key_masked_when_passed_from_env(
-    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")
-    llm = OpenAI(openai_api_key="secret-api-key")
-    print(llm.openai_api_key, end="")  # noqa: T201
-    captured = capsys.readouterr()
-
-    assert captured.out == "**********"
-
-
-def test_openai_api_key_masked_when_passed_via_constructor(
-    capsys: CaptureFixture,
-) -> None:
-    """Test that the API key is masked when passed via the constructor."""
-    llm = OpenAI(openai_api_key="secret-api-key")
-    print(llm.openai_api_key, end="")  # noqa: T201
-    captured = capsys.readouterr()
-
-    assert captured.out == "**********"
-
-
-def test_openai_uses_actual_secret_value_from_secretstr() -> None:
-    """Test that the actual secret value is correctly retrieved."""
-    llm = OpenAI(openai_api_key="secret-api-key")
-    assert cast(SecretStr, llm.openai_api_key).get_secret_value() == "secret-api-key"
-
-
-# openai embeddings
-def test_openai_embeddings_api_key_is_secret_string() -> None:
-    """Test that the API key is stored as a SecretStr."""
-    embeddings_model = OpenAIEmbeddings(openai_api_key="secret-api-key")
-    assert isinstance(embeddings_model.openai_api_key, SecretStr)
-
-
-def test_openai_embeddings_api_key_masked_when_passed_from_env(
-    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")
-    embeddings_model = OpenAIEmbeddings(openai_api_key="secret-api-key")
-    print(embeddings_model.openai_api_key, end="")  # noqa: T201
-    captured = capsys.readouterr()
-
-    assert captured.out == "**********"
-
-
-def test_openai_embeddings_api_key_masked_when_passed_via_constructor(
-    capsys: CaptureFixture,
-) -> None:
-    """Test that the API key is masked when passed via the constructor."""
-    embeddings_model = OpenAIEmbeddings(openai_api_key="secret-api-key")
-    print(embeddings_model.openai_api_key, end="")  # noqa: T201
-    captured = capsys.readouterr()
-
-    assert captured.out == "**********"
-
-
-def test_openai_embeddings_uses_actual_secret_value_from_secretstr() -> None:
-    """Test that the actual secret value is correctly retrieved."""
-    embeddings_model = OpenAIEmbeddings(openai_api_key="secret-api-key")
-    assert (
-        cast(SecretStr, embeddings_model.openai_api_key).get_secret_value()
-        == "secret-api-key"
-    )