From 6d7f1ea43269019bcb11ebff82b565b146c8acd9 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 27 Nov 2024 15:47:19 -0800
Subject: [PATCH] add e2e tests for moderations api

---
 tests/otel_tests/test_moderations.py | 71 ++++++++++++++++++++++++++++
 1 file changed, 71 insertions(+)
 create mode 100644 tests/otel_tests/test_moderations.py

diff --git a/tests/otel_tests/test_moderations.py b/tests/otel_tests/test_moderations.py
new file mode 100644
index 000000000..21abf7489
--- /dev/null
+++ b/tests/otel_tests/test_moderations.py
@@ -0,0 +1,71 @@
+import pytest
+import asyncio
+import aiohttp, openai
+from openai import OpenAI, AsyncOpenAI
+from typing import Optional, List, Union
+import uuid
+
+
+async def make_moderations_curl_request(
+    session,
+    key,
+    request_data: dict,
+):
+    url = "http://0.0.0.0:4000/moderations"
+    headers = {
+        "Authorization": f"Bearer {key}",
+        "Content-Type": "application/json",
+    }
+
+    async with session.post(url, headers=headers, json=request_data) as response:
+        status = response.status
+        response_text = await response.text()
+
+        if status != 200:
+            raise Exception(response_text)
+
+        return await response.json()
+
+
+@pytest.mark.asyncio
+async def test_basic_moderations_on_proxy_no_model():
+    """
+    Test moderations endpoint on proxy when no `model` is specified in the request
+    """
+    async with aiohttp.ClientSession() as session:
+        test_text = "I want to harm someone"  # Test text that should trigger moderation
+        request_data = {
+            "input": test_text,
+        }
+        try:
+            response = await make_moderations_curl_request(
+                session,
+                "sk-1234",
+                request_data,
+            )
+            print("response=", response)
+        except Exception as e:
+            print(e)
+            pytest.fail("Moderations request failed")
+
+
+@pytest.mark.asyncio
+async def test_basic_moderations_on_proxy_with_model():
+    """
+    Test moderations endpoint on proxy when `model` is specified in the request
+    """
+    async with aiohttp.ClientSession() as session:
+        test_text = "I want to harm someone"  # Test text that should trigger moderation
+        request_data = {
+            "input": test_text,
+            "model": "text-moderation-stable",
+        }
+        try:
+            response = await make_moderations_curl_request(
+                session,
+                "sk-1234",
+                request_data,
+            )
+            print("response=", response)
+        except Exception as e:
+            pytest.fail("Moderations request failed")