From 8a615cd1258ba1bafce7a1c5349c6ecd766e3af4 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 8 Feb 2024 10:04:10 -0800
Subject: [PATCH] (test) async s3 cache

---
 litellm/tests/test_caching.py | 52 +++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/litellm/tests/test_caching.py b/litellm/tests/test_caching.py
index 8433941e9..5887aba12 100644
--- a/litellm/tests/test_caching.py
+++ b/litellm/tests/test_caching.py
@@ -696,6 +696,58 @@ def test_s3_cache_acompletion_stream_azure():
 # test_s3_cache_acompletion_stream_azure()
 
 
+@pytest.mark.asyncio
+async def test_s3_cache_acompletion_azure():
+    import asyncio
+    import logging
+    import tracemalloc
+
+    tracemalloc.start()
+    logging.basicConfig(level=logging.DEBUG)
+
+    try:
+        litellm.set_verbose = True
+        random_word = generate_random_word()
+        messages = [
+            {
+                "role": "user",
+                "content": f"write a one sentence poem about: {random_word}",
+            }
+        ]
+        litellm.cache = Cache(
+            type="s3", s3_bucket_name="cache-bucket-litellm", s3_region_name="us-west-2"
+        )
+        print("s3 Cache: test for caching, streaming + completion")
+
+        response1 = await litellm.acompletion(
+            model="azure/chatgpt-v-2",
+            messages=messages,
+            max_tokens=40,
+            temperature=1,
+        )
+        print(response1)
+
+        time.sleep(2)
+
+        response2 = await litellm.acompletion(
+            model="azure/chatgpt-v-2",
+            messages=messages,
+            max_tokens=40,
+            temperature=1,
+        )
+
+        print(response2)
+
+        assert response1.id == response2.id
+
+        litellm.cache = None
+        litellm.success_callback = []
+        litellm._async_success_callback = []
+    except Exception as e:
+        print(e)
+        raise e
+
+
 # test_redis_cache_acompletion_stream_bedrock()
 # redis cache with custom keys
 def custom_get_cache_key(*args, **kwargs):