From 15b1e758ab4e82ec1c7c0c8f2e6e25e2f756d22b Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 15 Dec 2023 16:43:33 +0530
Subject: [PATCH] (test) dynamo db writes

---
 .gitignore                          |  1 +
 litellm/integrations/dynamodb.py    |  5 ++-
 litellm/tests/test_dynamodb_logs.py | 68 ++++++++++++++++++++++++++++-
 3 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/.gitignore b/.gitignore
index d2be5dee2..c5ad13111 100644
--- a/.gitignore
+++ b/.gitignore
@@ -22,4 +22,5 @@ litellm/tests/config_*.yaml
 litellm/tests/langfuse.log
 litellm/tests/test_custom_logger.py
 litellm/tests/langfuse.log
+litellm/tests/dynamo.log
 .vscode/settings.json
diff --git a/litellm/integrations/dynamodb.py b/litellm/integrations/dynamodb.py
index cb9a32098..7f7111fd9 100644
--- a/litellm/integrations/dynamodb.py
+++ b/litellm/integrations/dynamodb.py
@@ -27,6 +27,7 @@ class DyanmoDBLogger:
 
     def check_table_exists(self):
         existing_tables = self.dynamodb.meta.client.list_tables()['TableNames']
+        print_verbose(f"Dynamo DB: Existing Tables= {existing_tables}")
         return self.table_name in existing_tables
 
 
@@ -66,7 +67,6 @@ class DyanmoDBLogger:
             id = response_obj.get("id", str(uuid.uuid4()))
 
  
-
             # convert all optional params to str
             for param, value in optional_params.items():
                 try:
@@ -95,7 +95,8 @@ class DyanmoDBLogger:
             table = self.dynamodb.Table(self.table_name)
             # Assuming log_data is a dictionary with log information
             response = table.put_item(Item=payload)
-            print(f'Log data added to {self.table_name} successfully:', response)
+
+            print_verbose(f"Response from DynamoDB:{str(response)}")
 
             print_verbose(
                 f"DynamoDB Layer Logging - final response object: {response_obj}"
diff --git a/litellm/tests/test_dynamodb_logs.py b/litellm/tests/test_dynamodb_logs.py
index 82016cefc..ca2345189 100644
--- a/litellm/tests/test_dynamodb_logs.py
+++ b/litellm/tests/test_dynamodb_logs.py
@@ -10,12 +10,66 @@ import litellm
 litellm.num_retries = 3
 litellm.success_callback = ["dynamodb"]
 
+litellm.set_verbose = True
+
+
 import time
 import pytest
 
+def verify_dynamo_logs():
+    num_requests = 2
+    pass
+
+
+def pre_request():
+    log_file = open("dynamo.log", "a+")
+
+    # Clear the contents of the file by truncating it
+    log_file.truncate(0)
+
+    # Save the original stdout so that we can restore it later
+    original_stdout = sys.stdout
+    # Redirect stdout to the file
+    sys.stdout = log_file
+
+    return original_stdout, log_file
+
+
+import re
+def verify_log_file(log_file_path):
+
+    with open(log_file_path, 'r') as log_file:
+        log_content = log_file.read()
+
+        # Define the pattern to search for in the log file
+        pattern = r"Response from DynamoDB:{.*?}"
+
+        # Find all matches in the log content
+        matches = re.findall(pattern, log_content)
+
+        # Print the DynamoDB success log matches
+        print("DynamoDB Success Log Matches:")
+        for match in matches:
+            print(match)
+
+        # Print the total count of lines containing the specified response
+        print(f"Total occurrences of specified response: {len(matches)}")
+
+        # Count the occurrences of successful responses (status code 200 or 201)
+        success_count = sum(1 for match in matches if "'HTTPStatusCode': 200" in match or "'HTTPStatusCode': 201" in match)
+
+        # Print the count of successful responses
+        print(f"Count of successful responses from DynamoDB: {success_count}")
+    assert success_count == 5
+
+   
+
+
+
 def test_dynamo_logging_async(): 
     try: 
-        litellm.set_verbose = True
+        # pre
+        original_stdout, log_file = pre_request()
         async def _test():
             return await litellm.acompletion(
                 model="gpt-3.5-turbo",
@@ -26,12 +80,22 @@ def test_dynamo_logging_async():
             )
         response = asyncio.run(_test())
         print(f"response: {response}")
+        time.sleep(1)
     except litellm.Timeout as e: 
         pass
     except Exception as e: 
         pytest.fail(f"An exception occurred - {e}")
+    finally:
+        # post, close log file and verify
+        # Reset stdout to the original value
+        sys.stdout = original_stdout
+        # Close the file
+        log_file.close()
+        verify_log_file("dynamo.log")
 
-# test_dynamo_logging_async()
+
+
+test_dynamo_logging_async()
 
 
 def test_dynamo_logging_async_stream():