add test_no_async_http_handler_usage

This commit is contained in:
Ishaan Jaff 2024-11-21 10:16:07 -08:00
parent 77232f9bc4
commit 4d56249eb9

View file

@ -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()