From fe6e46546649f7368510ae0ba984e83cb16c75ae Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 3 May 2024 09:59:59 -0700
Subject: [PATCH] test - num callbacks on proxy should not increase

---
 tests/test_callbacks_on_proxy.py | 79 ++++++++++++++++++++++++++++++--
 1 file changed, 75 insertions(+), 4 deletions(-)

diff --git a/tests/test_callbacks_on_proxy.py b/tests/test_callbacks_on_proxy.py
index 6effc7464..01183cf6e 100644
--- a/tests/test_callbacks_on_proxy.py
+++ b/tests/test_callbacks_on_proxy.py
@@ -16,12 +16,13 @@ import pytest
 load_dotenv()
 
 
-async def config_update(session):
+async def config_update(session, routing_strategy=None):
     url = "http://0.0.0.0:4000/config/update"
     headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+    print("routing_strategy: ", routing_strategy)
     data = {
         "router_settings": {
-            "routing_strategy": ["latency-based-routing"],
+            "routing_strategy": routing_strategy,
         },
     }
 
@@ -60,9 +61,41 @@ async def get_active_callbacks(session):
         return _num_callbacks
 
 
+async def get_current_routing_strategy(session):
+    url = "http://0.0.0.0:4000/get/config/callbacks"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": "Bearer sk-1234",
+    }
+
+    async with session.get(url, headers=headers) as response:
+        status = response.status
+        response_text = await response.text()
+        print(response_text)
+        print()
+
+        if status != 200:
+            raise Exception(f"Request did not return a 200 status code: {status}")
+
+        _json_response = await response.json()
+        print("JSON response: ", _json_response)
+
+        router_settings = _json_response["router_settings"]
+        print("Router settings: ", router_settings)
+        routing_strategy = router_settings["routing_strategy"]
+        return routing_strategy
+
+
 @pytest.mark.asyncio
-async def test_add_model_run_health():
-    """ """
+async def test_check_num_callbacks():
+    """
+    Test 1:  num callbacks should NOT increase over time
+    -> check current callbacks
+    -> sleep for 30s
+    -> check current callbacks
+    -> sleep for 30s
+    -> check current callbacks
+    """
     import uuid
 
     async with aiohttp.ClientSession() as session:
@@ -72,8 +105,46 @@ async def test_add_model_run_health():
 
         num_callbacks_2 = await get_active_callbacks(session=session)
 
+        assert num_callbacks_1 == num_callbacks_2
+
         await asyncio.sleep(30)
 
         num_callbacks_3 = await get_active_callbacks(session=session)
 
         assert num_callbacks_1 == num_callbacks_2 == num_callbacks_3
+
+
+@pytest.mark.asyncio
+async def test_check_num_callbacks_on_lowest_latency():
+    """
+    Test 1:  num callbacks should NOT increase over time
+    -> Update to lowest latency
+    -> check current callbacks
+    -> sleep for 30s
+    -> check current callbacks
+    -> sleep for 30s
+    -> check current callbacks
+    -> update back to original routing-strategy
+    """
+    import uuid
+
+    async with aiohttp.ClientSession() as session:
+
+        original_routing_strategy = await get_current_routing_strategy(session=session)
+        await config_update(session=session, routing_strategy="latency-based-routing")
+
+        num_callbacks_1 = await get_active_callbacks(session=session)
+
+        await asyncio.sleep(30)
+
+        num_callbacks_2 = await get_active_callbacks(session=session)
+
+        assert num_callbacks_1 == num_callbacks_2
+
+        await asyncio.sleep(30)
+
+        num_callbacks_3 = await get_active_callbacks(session=session)
+
+        assert num_callbacks_1 == num_callbacks_2 == num_callbacks_3
+
+        await config_update(session=session, routing_strategy=original_routing_strategy)