From 3a4e3e6ac0aeacb0d9e0ee6498df73e91a44d9bb Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 13 Nov 2024 14:38:20 -0800
Subject: [PATCH] test_get_secret_with_access_mode

---
 litellm/proxy/proxy_config.yaml            |  3 +-
 litellm/secret_managers/main.py            | 11 ++--
 tests/local_testing/test_secret_manager.py | 75 +++++++++++++++++++++-
 3 files changed, 81 insertions(+), 8 deletions(-)

diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index d36ef5df0..71e3dee0e 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -10,4 +10,5 @@ model_list:
 general_settings:
   key_management_system: "aws_secret_manager" 
   key_management_settings: 
-    store_virtual_keys: true
\ No newline at end of file
+    store_virtual_keys: true
+    access_mode: "write_only"
diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py
index 5c47fbd00..ce6d30755 100644
--- a/litellm/secret_managers/main.py
+++ b/litellm/secret_managers/main.py
@@ -335,9 +335,10 @@ def _should_read_secret_from_secret_manager() -> bool:
     - Otherwise, return False
     """
     if litellm.secret_manager_client is not None:
-        if (
-            litellm._key_management_settings.access_mode == "read_only"
-            or litellm._key_management_settings.access_mode == "read_and_write"
-        ):
-            return True
+        if litellm._key_management_settings is not None:
+            if (
+                litellm._key_management_settings.access_mode == "read_only"
+                or litellm._key_management_settings.access_mode == "read_and_write"
+            ):
+                return True
     return False
diff --git a/tests/local_testing/test_secret_manager.py b/tests/local_testing/test_secret_manager.py
index 698812fbb..1b95119a3 100644
--- a/tests/local_testing/test_secret_manager.py
+++ b/tests/local_testing/test_secret_manager.py
@@ -15,11 +15,14 @@ sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
 import pytest
-
+import litellm
 from litellm.llms.AzureOpenAI.azure import get_azure_ad_token_from_oidc
 from litellm.llms.bedrock.chat import BedrockConverseLLM, BedrockLLM
 from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
-from litellm.secret_managers.main import get_secret
+from litellm.secret_managers.main import (
+    get_secret,
+    _should_read_secret_from_secret_manager,
+)
 
 
 def test_aws_secret_manager():
@@ -244,3 +247,71 @@ def test_google_secret_manager_read_in_memory():
     )
     print("secret_val: {}".format(secret_val))
     assert secret_val == "lite-llm"
+
+
+def test_should_read_secret_from_secret_manager():
+    """
+    Test that _should_read_secret_from_secret_manager returns correct values based on access mode
+    """
+    from litellm.proxy._types import KeyManagementSettings
+
+    # Test when secret manager client is None
+    litellm.secret_manager_client = None
+    litellm._key_management_settings = KeyManagementSettings()
+    assert _should_read_secret_from_secret_manager() is False
+
+    # Test with secret manager client and read_only access
+    litellm.secret_manager_client = "dummy_client"
+    litellm._key_management_settings = KeyManagementSettings(access_mode="read_only")
+    assert _should_read_secret_from_secret_manager() is True
+
+    # Test with secret manager client and read_and_write access
+    litellm._key_management_settings = KeyManagementSettings(
+        access_mode="read_and_write"
+    )
+    assert _should_read_secret_from_secret_manager() is True
+
+    # Test with secret manager client and write_only access
+    litellm._key_management_settings = KeyManagementSettings(access_mode="write_only")
+    assert _should_read_secret_from_secret_manager() is False
+
+    # Reset global variables
+    litellm.secret_manager_client = None
+    litellm._key_management_settings = KeyManagementSettings()
+
+
+def test_get_secret_with_access_mode():
+    """
+    Test that get_secret respects access mode settings
+    """
+    from litellm.proxy._types import KeyManagementSettings
+
+    # Set up test environment
+    test_secret_name = "TEST_SECRET_KEY"
+    test_secret_value = "test_secret_value"
+    os.environ[test_secret_name] = test_secret_value
+
+    # Test with write_only access (should read from os.environ)
+    litellm.secret_manager_client = "dummy_client"
+    litellm._key_management_settings = KeyManagementSettings(access_mode="write_only")
+    assert get_secret(test_secret_name) == test_secret_value
+
+    # Test with no KeyManagementSettings but secret_manager_client set
+    litellm.secret_manager_client = "dummy_client"
+    litellm._key_management_settings = KeyManagementSettings()
+    assert _should_read_secret_from_secret_manager() is True
+
+    # Test with read_only access
+    litellm._key_management_settings = KeyManagementSettings(access_mode="read_only")
+    assert _should_read_secret_from_secret_manager() is True
+
+    # Test with read_and_write access
+    litellm._key_management_settings = KeyManagementSettings(
+        access_mode="read_and_write"
+    )
+    assert _should_read_secret_from_secret_manager() is True
+
+    # Reset global variables
+    litellm.secret_manager_client = None
+    litellm._key_management_settings = KeyManagementSettings()
+    del os.environ[test_secret_name]