From 4d56249eb97423a19c59f93bf79c592ccd29fee7 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 21 Nov 2024 10:16:07 -0800
Subject: [PATCH] add test_no_async_http_handler_usage

---
 .../ensure_async_clients_test.py              | 86 +++++++++++++++++++
 1 file changed, 86 insertions(+)
 create mode 100644 tests/code_coverage_tests/ensure_async_clients_test.py

diff --git a/tests/code_coverage_tests/ensure_async_clients_test.py b/tests/code_coverage_tests/ensure_async_clients_test.py
new file mode 100644
index 000000000..d65d56f64
--- /dev/null
+++ b/tests/code_coverage_tests/ensure_async_clients_test.py
@@ -0,0 +1,86 @@
+import ast
+import os
+
+ALLOWED_FILES = [
+    # local files
+    "../../litellm/__init__.py",
+    "../../litellm/llms/custom_httpx/http_handler.py",
+    # when running on ci/cd
+    "./litellm/__init__.py",
+    "./litellm/llms/custom_httpx/http_handler.py",
+]
+
+
+def check_for_async_http_handler(file_path):
+    """
+    Checks if AsyncHttpHandler is instantiated in the given file.
+    Returns a list of line numbers where AsyncHttpHandler is used.
+    """
+    print("..checking file=", file_path)
+    if file_path in ALLOWED_FILES:
+        return []
+    with open(file_path, "r") as file:
+        try:
+            tree = ast.parse(file.read())
+        except SyntaxError:
+            print(f"Warning: Syntax error in file {file_path}")
+            return []
+
+    violations = []
+    target_names = [
+        "AsyncHttpHandler",
+        "AsyncHTTPHandler",
+        "AsyncClient",
+        "httpx.AsyncClient",
+    ]  # Add variations here
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Call):
+            if isinstance(node.func, ast.Name) and node.func.id.lower() in [
+                name.lower() for name in target_names
+            ]:
+                raise ValueError(
+                    f"found violation in file {file_path} line: {node.lineno}"
+                )
+    return violations
+
+
+def scan_directory_for_async_handler(base_dir):
+    """
+    Scans all Python files in the directory tree for AsyncHttpHandler usage.
+    Returns a dict of files and line numbers where violations were found.
+    """
+    violations = {}
+
+    for root, _, files in os.walk(base_dir):
+        for file in files:
+            if file.endswith(".py"):
+                file_path = os.path.join(root, file)
+                file_violations = check_for_async_http_handler(file_path)
+                if file_violations:
+                    violations[file_path] = file_violations
+
+    return violations
+
+
+def test_no_async_http_handler_usage():
+    """
+    Test to ensure AsyncHttpHandler is not used anywhere in the codebase.
+    """
+    base_dir = "./litellm"  # Adjust this path as needed
+
+    # base_dir = "../../litellm"  # LOCAL TESTING
+    violations = scan_directory_for_async_handler(base_dir)
+
+    if violations:
+        violation_messages = []
+        for file_path, line_numbers in violations.items():
+            violation_messages.append(
+                f"Found AsyncHttpHandler in {file_path} at lines: {line_numbers}"
+            )
+        raise AssertionError(
+            "AsyncHttpHandler usage detected:\n" + "\n".join(violation_messages)
+        )
+
+
+if __name__ == "__main__":
+    test_no_async_http_handler_usage()