From 4dab664cdc86f1cc6f0ec8eb095ac82dfd1cf1f4 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 13 Nov 2024 11:54:00 -0800
Subject: [PATCH] add test for async_delete_secret

---
 .../local_testing/test_aws_secret_manager.py  | 78 ++++++++++++-------
 1 file changed, 49 insertions(+), 29 deletions(-)

diff --git a/tests/local_testing/test_aws_secret_manager.py b/tests/local_testing/test_aws_secret_manager.py
index 45ee13c33..baa237dfd 100644
--- a/tests/local_testing/test_aws_secret_manager.py
+++ b/tests/local_testing/test_aws_secret_manager.py
@@ -50,26 +50,36 @@ async def test_write_and_read_simple_secret():
     test_secret_name = f"litellm_test_{uuid.uuid4().hex[:8]}"
     test_secret_value = "test_value_123"
 
-    # Write secret
-    write_response = await secret_manager.async_write_secret(
-        secret_name=test_secret_name,
-        secret_value=test_secret_value,
-        description="LiteLLM Test Secret",
-    )
+    try:
+        # Write secret
+        write_response = await secret_manager.async_write_secret(
+            secret_name=test_secret_name,
+            secret_value=test_secret_value,
+            description="LiteLLM Test Secret",
+        )
 
-    print("Write Response:", write_response)
+        print("Write Response:", write_response)
 
-    assert write_response is not None
-    assert "ARN" in write_response
-    assert "Name" in write_response
-    assert write_response["Name"] == test_secret_name
+        assert write_response is not None
+        assert "ARN" in write_response
+        assert "Name" in write_response
+        assert write_response["Name"] == test_secret_name
 
-    # Read secret back
-    read_value = await secret_manager.async_read_secret(secret_name=test_secret_name)
+        # Read secret back
+        read_value = await secret_manager.async_read_secret(
+            secret_name=test_secret_name
+        )
 
-    print("Read Value:", read_value)
+        print("Read Value:", read_value)
 
-    assert read_value == test_secret_value
+        assert read_value == test_secret_value
+    finally:
+        # Cleanup: Delete the secret
+        delete_response = await secret_manager.async_delete_secret(
+            secret_name=test_secret_name
+        )
+        print("Delete Response:", delete_response)
+        assert delete_response is not None
 
 
 @pytest.mark.asyncio
@@ -86,24 +96,34 @@ async def test_write_and_read_json_secret():
         "metadata": {"team": "ml", "project": "litellm"},
     }
 
-    # Write JSON secret
-    write_response = await secret_manager.async_write_secret(
-        secret_name=test_secret_name,
-        secret_value=json.dumps(test_secret_value),
-        description="LiteLLM JSON Test Secret",
-    )
+    try:
+        # Write JSON secret
+        write_response = await secret_manager.async_write_secret(
+            secret_name=test_secret_name,
+            secret_value=json.dumps(test_secret_value),
+            description="LiteLLM JSON Test Secret",
+        )
 
-    print("Write Response:", write_response)
+        print("Write Response:", write_response)
 
-    # Read and parse JSON secret
-    read_value = await secret_manager.async_read_secret(secret_name=test_secret_name)
-    parsed_value = json.loads(read_value)
+        # Read and parse JSON secret
+        read_value = await secret_manager.async_read_secret(
+            secret_name=test_secret_name
+        )
+        parsed_value = json.loads(read_value)
 
-    print("Read Value:", read_value)
+        print("Read Value:", read_value)
 
-    assert parsed_value == test_secret_value
-    assert parsed_value["api_key"] == "test_key"
-    assert parsed_value["metadata"]["team"] == "ml"
+        assert parsed_value == test_secret_value
+        assert parsed_value["api_key"] == "test_key"
+        assert parsed_value["metadata"]["team"] == "ml"
+    finally:
+        # Cleanup: Delete the secret
+        delete_response = await secret_manager.async_delete_secret(
+            secret_name=test_secret_name
+        )
+        print("Delete Response:", delete_response)
+        assert delete_response is not None
 
 
 @pytest.mark.asyncio