diff --git a/litellm/tests/test_s3_logs.py b/litellm/tests/test_s3_logs.py index a5347fb084d..b2c9bd0d279 100644 --- a/litellm/tests/test_s3_logs.py +++ b/litellm/tests/test_s3_logs.py @@ -115,4 +115,99 @@ def test_s3_logging(): print("Passed! Testing async s3 logging") -test_s3_logging() +# test_s3_logging() + + +def test_s3_logging_r2(): + # all s3 requests need to be in one test function + # since we are modifying stdout, and pytests runs tests in parallel + # on circle ci - we only test litellm.acompletion() + try: + # redirect stdout to log_file + # litellm.cache = litellm.Cache( + # type="s3", s3_bucket_name="litellm-r2-bucket", s3_region_name="us-west-2" + # ) + litellm.set_verbose = True + + litellm.success_callback = ["s3"] + litellm.s3_callback_params = { + "s3_bucket_name": "litellm-r2-bucket", + "s3_aws_secret_access_key": "os.environ/R2_S3_ACCESS_KEY", + "s3_aws_access_key_id": "os.environ/R2_S3_ACCESS_ID", + "s3_endpoint_url": "os.environ/R2_S3_URL", + "s3_region_name": "os.environ/R2_S3_REGION_NAME", + } + print("Testing async s3 logging") + + expected_keys = [] + + import time + + curr_time = str(time.time()) + + async def _test(): + return await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": f"This is a test {curr_time}"}], + max_tokens=10, + temperature=0.7, + user="ishaan-2", + ) + + response = asyncio.run(_test()) + print(f"response: {response}") + expected_keys.append(response.id) + + import boto3 + + s3 = boto3.client( + "s3", + endpoint_url="https://0399b10e77ac6668c80404a5ff49eb37.r2.cloudflarestorage.com", + region_name="us-east-1", + aws_access_key_id="c93b3b5e490a18391edd8c05ffe8eb9c", + aws_secret_access_key="c8b76230569cb6318e7c30959432b1e03f08a86cc1641e873816eb6eadd6e694", + ) + + bucket_name = "litellm-r2-bucket" + # List objects in the bucket + response = s3.list_objects(Bucket=bucket_name) + + # Sort the objects based on the LastModified timestamp + objects = sorted( + response["Contents"], key=lambda x: x["LastModified"], reverse=True + ) + # Get the keys of the most recent objects + most_recent_keys = [obj["Key"] for obj in objects] + print(most_recent_keys) + # for each key, get the part before "-" as the key. Do it safely + cleaned_keys = [] + for key in most_recent_keys: + split_key = key.split("-time=") + cleaned_keys.append(split_key[0]) + print("\n most recent keys", most_recent_keys) + print("\n cleaned keys", cleaned_keys) + print("\n Expected keys: ", expected_keys) + matches = 0 + for key in expected_keys: + assert key in cleaned_keys + + if key in cleaned_keys: + matches += 1 + # remove the match key + cleaned_keys.remove(key) + # this asserts we log, the first request + the 2nd cached request + print("we had two matches ! passed ", matches) + assert matches == 1 + try: + # cleanup s3 bucket in test + for key in most_recent_keys: + s3.delete_object(Bucket=bucket_name, Key=key) + except: + # don't let cleanup fail a test + 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 + print("Passed! Testing async s3 logging")