mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-04 23:11:19 +00:00
Add end to end integration tests for batches
This commit is contained in:
Ephrim Stanley
@@ -1,221 +0,0 @@
|
||||
"""Base class for managed files and batch API tests."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
from tenacity import Retrying, stop_after_delay, wait_fixed
|
||||
|
||||
|
||||
LOCAL_LITELLM_BASE_URL = "http://localhost:4000"
|
||||
LOCAL_AZURE_BASE_URL = "http://localhost:8090"
|
||||
|
||||
USE_LITELLM = os.environ.get("USE_LITELLM", "true").lower() == "true"
|
||||
if USE_LITELLM:
|
||||
base_url = LOCAL_LITELLM_BASE_URL
|
||||
api_key = "sk-1234"
|
||||
else:
|
||||
base_url = LOCAL_AZURE_BASE_URL
|
||||
api_key = "sk-1234"
|
||||
|
||||
USE_MOCK_SERVER = os.environ.get("USE_MOCK_SERVER", "false").lower() == "true"
|
||||
if USE_MOCK_SERVER:
|
||||
model_name = "azure-fake-gpt-5-batch-2025-08-07"
|
||||
MODEL_NAMES = [
|
||||
"azure-fake-gpt-5-batch-2025-08-07",
|
||||
# "anthropic-fake-claude-sonnet-4-batch-2025-08-07",
|
||||
# "vertex-fake-gemini-2.5-pro-batch-2025-08-07",
|
||||
]
|
||||
else:
|
||||
model_name = "gpt-5-batch-2025-08-07"
|
||||
MODEL_NAMES = [
|
||||
"gpt-5-batch-2025-08-07",
|
||||
# "claude-sonnet-4-batch-2025-08-07",
|
||||
# "gemini-2.5-pro-batch-2025-08-07",
|
||||
]
|
||||
|
||||
|
||||
def _extract_model_id(model_name: str) -> str:
|
||||
if "gpt" in model_name:
|
||||
return "gpt"
|
||||
elif "claude" in model_name or "anthropic" in model_name:
|
||||
return "anthropic"
|
||||
elif "gemini" in model_name or "vertex" in model_name:
|
||||
return "gemini"
|
||||
return model_name.split("-")[0]
|
||||
|
||||
|
||||
MODEL_IDS = [_extract_model_id(m) for m in MODEL_NAMES]
|
||||
|
||||
MIN_EXPIRY_SECONDS = 259200
|
||||
|
||||
|
||||
class ManagedFilesBase:
|
||||
"""Base class with shared helpers for managed files and batch tests."""
|
||||
|
||||
base_url = base_url
|
||||
api_key = api_key
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup_test(self):
|
||||
print(f"Base URL: {self.base_url}, Model: {model_name}\n")
|
||||
self.reset_mock_server()
|
||||
|
||||
@staticmethod
|
||||
def generate_request_id():
|
||||
return f"req-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
def create_openai_client(self, api_key: str) -> openai.OpenAI:
|
||||
return openai.OpenAI(
|
||||
base_url=self.base_url,
|
||||
api_key=api_key,
|
||||
http_client=httpx.Client(verify=False),
|
||||
)
|
||||
|
||||
def create_batch_request_file_on_disk(self, tmpdir, model: str):
|
||||
request_id = self.generate_request_id()
|
||||
batch_request = {
|
||||
"custom_id": request_id,
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
request_file = os.path.join(tmpdir, f"request-{request_id}.jsonl")
|
||||
with open(request_file, "w") as f:
|
||||
f.write(json.dumps(batch_request))
|
||||
|
||||
return request_file
|
||||
|
||||
def create_batch_input_file(
|
||||
self,
|
||||
client: openai.OpenAI,
|
||||
request_file: str,
|
||||
expiry_seconds: int = MIN_EXPIRY_SECONDS,
|
||||
):
|
||||
batch_input_file = client.files.create(
|
||||
file=open(request_file, "rb"),
|
||||
purpose="batch",
|
||||
extra_body={
|
||||
"target_model_names": model_name,
|
||||
"expires_after": {
|
||||
"seconds": expiry_seconds,
|
||||
"anchor": "created_at",
|
||||
},
|
||||
},
|
||||
)
|
||||
return batch_input_file
|
||||
|
||||
def create_batch(
|
||||
self,
|
||||
client: openai.OpenAI,
|
||||
input_file_id: str,
|
||||
expiry_seconds: int = MIN_EXPIRY_SECONDS,
|
||||
):
|
||||
batch = client.batches.create(
|
||||
input_file_id=input_file_id,
|
||||
endpoint="/v1/chat/completions",
|
||||
completion_window="24h",
|
||||
extra_body={
|
||||
"output_expires_after": {
|
||||
"seconds": expiry_seconds,
|
||||
"anchor": "created_at",
|
||||
},
|
||||
},
|
||||
)
|
||||
return batch
|
||||
|
||||
def wait_for_batch_state(
|
||||
self,
|
||||
client: openai.OpenAI,
|
||||
batch_id: str,
|
||||
expected_status: str,
|
||||
max_seconds: int = 60,
|
||||
wait_seconds: int = 5,
|
||||
):
|
||||
for attempt in Retrying(
|
||||
stop=stop_after_delay(max_seconds),
|
||||
wait=wait_fixed(wait_seconds),
|
||||
):
|
||||
with attempt:
|
||||
batch_response = client.batches.retrieve(batch_id=batch_id)
|
||||
print(
|
||||
f"[{time.strftime('%H:%M:%S')}] Batch status: {batch_response.status}, expected: {expected_status}",
|
||||
)
|
||||
if batch_response.status == expected_status:
|
||||
return batch_response
|
||||
if batch_response.status in ["failed", "expired", "cancelled"]:
|
||||
raise Exception(
|
||||
f"Batch failed with status: {batch_response.status}",
|
||||
)
|
||||
raise Exception(f"Batch not in {expected_status} state yet")
|
||||
return None
|
||||
|
||||
def wait_for_batch_completed(
|
||||
self,
|
||||
client: openai.OpenAI,
|
||||
batch_id: str,
|
||||
max_seconds: int = 120,
|
||||
wait_seconds: int = 5,
|
||||
):
|
||||
return self.wait_for_batch_state(
|
||||
client,
|
||||
batch_id,
|
||||
"completed",
|
||||
max_seconds,
|
||||
wait_seconds,
|
||||
)
|
||||
|
||||
def shorten_id(self, id_str: str) -> str:
|
||||
if id_str is None:
|
||||
return "None"
|
||||
if len(id_str) <= 20:
|
||||
return id_str
|
||||
return id_str[:8] + "..." + id_str[-8:]
|
||||
|
||||
def reset_mock_server(self):
|
||||
if not USE_MOCK_SERVER:
|
||||
return
|
||||
print("Resetting mock server state...")
|
||||
reset_response = httpx.post(f"{LOCAL_AZURE_BASE_URL}/reset")
|
||||
assert reset_response.status_code == 200, f"Reset failed: {reset_response.text}"
|
||||
|
||||
def print_file_metadata(self, file_obj, label="File"):
|
||||
print(f"{label} metadata:")
|
||||
print(f"\tid={self.shorten_id(file_obj.id)}")
|
||||
print(f"\tobject={file_obj.object}")
|
||||
print(f"\tbytes={file_obj.bytes}")
|
||||
print(f"\tfilename={file_obj.filename}")
|
||||
print(f"\tpurpose={file_obj.purpose}")
|
||||
print(f"\tstatus={file_obj.status}")
|
||||
print(f"\tcreated_at={file_obj.created_at}")
|
||||
print(f"\texpires_at={file_obj.expires_at}")
|
||||
if file_obj.status_details:
|
||||
print(f"\tstatus_details={file_obj.status_details}")
|
||||
|
||||
def print_batch_metadata(self, batch):
|
||||
print("Batch metadata:")
|
||||
print(f"\tid={self.shorten_id(batch.id)}")
|
||||
print(f"\tstatus={batch.status}")
|
||||
print(f"\tendpoint={batch.endpoint}")
|
||||
print(f"\tcompletion_window={batch.completion_window}")
|
||||
print(f"\tinput_file_id={self.shorten_id(batch.input_file_id)}")
|
||||
print(f"\tcreated_at={batch.created_at}")
|
||||
print(f"\texpires_at={batch.expires_at}")
|
||||
print(f"\tin_progress_at={batch.in_progress_at}")
|
||||
print(f"\tcompleted_at={batch.completed_at}")
|
||||
print(f"\toutput_file_id={self.shorten_id(batch.output_file_id)}")
|
||||
print(f"\trequest_counts={batch.request_counts}")
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,205 +0,0 @@
|
||||
import warnings
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from tenacity import RetryError, Retrying, stop_after_delay, wait_fixed
|
||||
|
||||
from test_managed_files_base import (
|
||||
MODEL_IDS,
|
||||
MODEL_NAMES,
|
||||
ManagedFilesBase,
|
||||
MIN_EXPIRY_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
class TestManagedFilesAPI(ManagedFilesBase):
|
||||
"""Test cases for managed files and batch API.
|
||||
|
||||
Configuration via environment variables:
|
||||
USE_LITELLM=true - Run against LiteLLM proxy
|
||||
USE_LITELLM=false - Run against mock server directly (default)
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setup_class(cls):
|
||||
cls.openai_client = cls.create_openai_client(cls, cls.api_key)
|
||||
|
||||
def wait_for_batch_list(self, model_name, max_seconds=90, wait_seconds=10):
|
||||
for attempt in Retrying(
|
||||
stop=stop_after_delay(max_seconds),
|
||||
wait=wait_fixed(wait_seconds),
|
||||
):
|
||||
with attempt:
|
||||
batches_list = self.openai_client.batches.list(
|
||||
limit=10,
|
||||
extra_query={"target_model_names": model_name},
|
||||
)
|
||||
print(
|
||||
f"Batches in list: {len(batches_list.data)}",
|
||||
)
|
||||
if len(batches_list.data) == 0:
|
||||
raise Exception("No batches found in list yet")
|
||||
return batches_list
|
||||
return None
|
||||
|
||||
@pytest.mark.parametrize("model_name", MODEL_NAMES, ids=MODEL_IDS)
|
||||
def test_e2e_managed_batch(self, tmp_path, model_name):
|
||||
print(
|
||||
f"\n\nStarting test with base_url={self.base_url} and model_name={model_name}\n",
|
||||
)
|
||||
|
||||
self.reset_mock_server()
|
||||
|
||||
request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
|
||||
|
||||
print("Creating batch input file...")
|
||||
batch_input_file = self.create_batch_input_file(
|
||||
self.openai_client,
|
||||
request_file,
|
||||
MIN_EXPIRY_SECONDS,
|
||||
)
|
||||
print(f"Created batch input file: {self.shorten_id(batch_input_file.id)}\n")
|
||||
|
||||
print(
|
||||
f"Retrieving batch input file metadata for file id: {self.shorten_id(batch_input_file.id)}",
|
||||
)
|
||||
input_file_metadata = self.openai_client.files.retrieve(batch_input_file.id)
|
||||
assert input_file_metadata.id == batch_input_file.id, "File ID mismatch"
|
||||
assert input_file_metadata.object == "file", "object should be 'file'"
|
||||
assert input_file_metadata.bytes > 0, "bytes not set"
|
||||
assert input_file_metadata.filename == "modified_file.jsonl", (
|
||||
"filename mismatch"
|
||||
)
|
||||
assert input_file_metadata.purpose == "batch", "purpose mismatch"
|
||||
assert input_file_metadata.status in ["uploaded", "processed", "error"], (
|
||||
"invalid status"
|
||||
)
|
||||
assert input_file_metadata.created_at > 0, "created_at not set"
|
||||
if not input_file_metadata.expires_at:
|
||||
warnings.warn("batch input file expires_at not set")
|
||||
self.print_file_metadata(input_file_metadata, "Input file")
|
||||
|
||||
print("\nCreating batch...")
|
||||
batch = self.create_batch(
|
||||
self.openai_client,
|
||||
batch_input_file.id,
|
||||
MIN_EXPIRY_SECONDS,
|
||||
)
|
||||
print(f"Created batch: {self.shorten_id(batch.id)}")
|
||||
assert batch.id, "No batch ID returned"
|
||||
assert batch.input_file_id == batch_input_file.id, "File ID mismatch"
|
||||
assert batch.status in [
|
||||
"validating",
|
||||
"in_progress",
|
||||
"finalizing",
|
||||
"completed",
|
||||
], "Status mismatch"
|
||||
if not batch.expires_at:
|
||||
warnings.warn("batch expires_at not set")
|
||||
else:
|
||||
assert batch.expires_at > 0, "batch expires_at not set"
|
||||
|
||||
if not batch.endpoint:
|
||||
warnings.warn(
|
||||
"batch.endpoint empty in creation response - Azure API quirk, not a bug",
|
||||
)
|
||||
else:
|
||||
assert batch.endpoint == "/v1/chat/completions", "endpoint mismatch"
|
||||
assert batch.completion_window == "24h", "completion_window mismatch"
|
||||
assert batch.created_at > 0, "created_at not set"
|
||||
|
||||
self.print_batch_metadata(batch)
|
||||
|
||||
print("\nListing batches...")
|
||||
try:
|
||||
batches_list = self.wait_for_batch_list(
|
||||
model_name,
|
||||
max_seconds=30,
|
||||
wait_seconds=5,
|
||||
)
|
||||
batches = batches_list.data if batches_list else []
|
||||
batch_ids = [b.id for b in batches]
|
||||
if batch.id not in batch_ids:
|
||||
warnings.warn(
|
||||
f"Batch {batch.id} not found in list. batches.list returns raw IDs and not the encoded IDs. raw IDs: {batch_ids}",
|
||||
)
|
||||
except RetryError:
|
||||
warnings.warn(
|
||||
"batches.list() returned empty list after retries - known LiteLLM issue with managed batches",
|
||||
)
|
||||
except openai.APIError as e:
|
||||
pytest.fail(f"batches.list() failed: {e}")
|
||||
|
||||
print(
|
||||
f"\nWaiting for batch {self.shorten_id(batch.id)} to reach completed state...",
|
||||
)
|
||||
try:
|
||||
batch_response = self.wait_for_batch_state(
|
||||
self.openai_client,
|
||||
batch.id,
|
||||
"completed",
|
||||
max_seconds=30 * 60,
|
||||
wait_seconds=5,
|
||||
)
|
||||
except RetryError:
|
||||
raise TimeoutError("Timed out waiting for batch to be in state: completed")
|
||||
|
||||
print("\nRetrieving batch output file metadata...")
|
||||
output_file_metadata = self.openai_client.files.retrieve(
|
||||
batch_response.output_file_id,
|
||||
)
|
||||
assert output_file_metadata.id == batch_response.output_file_id, (
|
||||
"Output file ID mismatch"
|
||||
)
|
||||
assert output_file_metadata.object == "file", "object should be 'file'"
|
||||
assert output_file_metadata.bytes > 0, "bytes not set"
|
||||
assert output_file_metadata.filename, "filename not set"
|
||||
assert output_file_metadata.purpose in ["batch_output", "batch"], (
|
||||
"purpose should be batch_output"
|
||||
)
|
||||
assert output_file_metadata.created_at > 0, "created_at not set"
|
||||
self.print_file_metadata(output_file_metadata, "Output file")
|
||||
|
||||
print("\nFetching batch output file content...")
|
||||
batch_file_content = self.openai_client.files.content(
|
||||
batch_response.output_file_id,
|
||||
)
|
||||
assert batch_file_content.text, "No batch file content returned"
|
||||
assert len(batch_file_content.text) > 0, "Batch file content is empty"
|
||||
print(f"Output file content ({len(batch_file_content.text)} bytes):")
|
||||
for line in batch_file_content.text.strip().split("\n")[:3]:
|
||||
print(f"\t{line}")
|
||||
|
||||
print(f"\nDeleting input file: {self.shorten_id(batch_input_file.id)}")
|
||||
try:
|
||||
self.openai_client.files.delete(batch_input_file.id)
|
||||
except openai.APIError as e:
|
||||
pytest.fail(f"files.delete() failed: {e}")
|
||||
|
||||
print(
|
||||
f"\nDeleting output file: {self.shorten_id(batch_response.output_file_id)}",
|
||||
)
|
||||
try:
|
||||
self.openai_client.files.delete(batch_response.output_file_id)
|
||||
except openai.APIError as e:
|
||||
pytest.fail(f"files.delete() failed: {e}")
|
||||
|
||||
print("\nVerifying input file is deleted...")
|
||||
try:
|
||||
self.openai_client.files.content(batch_input_file.id)
|
||||
assert False, f"Input file {batch_input_file.id} exists after deletion"
|
||||
except (openai.NotFoundError, openai.PermissionDeniedError):
|
||||
print("Input file correctly not accessible after deletion")
|
||||
|
||||
print("\nVerifying output file is deleted...")
|
||||
try:
|
||||
self.openai_client.files.content(batch_response.output_file_id)
|
||||
assert False, (
|
||||
f"Output file {batch_response.output_file_id} exists after deletion"
|
||||
)
|
||||
except (openai.NotFoundError, openai.PermissionDeniedError):
|
||||
print("Output file correctly not accessible after deletion")
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,437 +0,0 @@
|
||||
"""
|
||||
Test cross-user batch access permissions.
|
||||
|
||||
This test verifies that a batch and related files created by one API key
|
||||
cannot be accessed, modified, or cancelled by a different API key.
|
||||
|
||||
Reference: https://github.com/BerriAI/litellm/pull/17401/files
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
|
||||
from test_managed_files_base import (
|
||||
ManagedFilesBase,
|
||||
MODEL_NAMES,
|
||||
MODEL_IDS,
|
||||
)
|
||||
|
||||
|
||||
BATCH_ROUTES = [
|
||||
"/v1/files",
|
||||
"/files",
|
||||
"/v1/files/*",
|
||||
"/files/*",
|
||||
"/v1/batches",
|
||||
"/batches",
|
||||
"/v1/batches/*",
|
||||
"/batches/*",
|
||||
]
|
||||
|
||||
|
||||
class TestManagedFilesPermissions(ManagedFilesBase):
|
||||
"""Test cases for cross-user batch access permissions.
|
||||
|
||||
Verifies that:
|
||||
- User A can create and access their own batches and files
|
||||
- User B cannot access, retrieve, cancel, or delete User A's batches/files
|
||||
"""
|
||||
|
||||
master_api_key = "sk-1234"
|
||||
|
||||
@classmethod
|
||||
def setup_class(cls):
|
||||
cls.admin_client = httpx.Client(base_url=cls.base_url, verify=False)
|
||||
|
||||
@classmethod
|
||||
def teardown_class(cls):
|
||||
cls.admin_client.close()
|
||||
|
||||
def user_suffix(self) -> str:
|
||||
return f"{time.strftime('%Y%m%d%H%M%S')}{int(time.time() * 1000) % 1000:03d}"
|
||||
|
||||
def create_user_and_key(self, user_suffix: str) -> tuple[str, str]:
|
||||
user_email = f"test-user-{user_suffix}-{self.user_suffix()}@test.com"
|
||||
|
||||
user_response = self.admin_client.post(
|
||||
"/user/new",
|
||||
json={
|
||||
"user_email": user_email,
|
||||
"user_alias": user_email,
|
||||
"user_role": "internal_user",
|
||||
"auto_create_key": "false",
|
||||
},
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.master_api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=30,
|
||||
)
|
||||
assert user_response.status_code == 200, (
|
||||
f"Failed to create user: {user_response.status_code} - {user_response.text}"
|
||||
)
|
||||
user_data = user_response.json()
|
||||
user_id = user_data.get("user_id")
|
||||
|
||||
key_response = self.admin_client.post(
|
||||
"/key/generate",
|
||||
json={
|
||||
"user_id": user_id,
|
||||
"key_alias": user_email,
|
||||
"allowed_routes": BATCH_ROUTES,
|
||||
},
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.master_api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=30,
|
||||
)
|
||||
assert key_response.status_code == 200, (
|
||||
f"Failed to create key: {key_response.status_code} - {key_response.text}"
|
||||
)
|
||||
key_data = key_response.json()
|
||||
api_key = key_data.get("key")
|
||||
|
||||
print(f"Created user {user_email} with key {key_data.get('key_alias')}")
|
||||
return user_id, api_key
|
||||
|
||||
def create_user_key_and_client(
|
||||
self,
|
||||
user_suffix: str,
|
||||
) -> tuple[str, str, openai.OpenAI]:
|
||||
user_id, api_key = self.create_user_and_key(user_suffix)
|
||||
return user_id, api_key, self.create_openai_client(api_key)
|
||||
|
||||
def create_key_and_client(self, user_id: str, key_suffix: str) -> str:
|
||||
key_alias = f"additional-key-{key_suffix}-{self.user_suffix()}"
|
||||
key_response = self.admin_client.post(
|
||||
"/key/generate",
|
||||
json={
|
||||
"user_id": user_id,
|
||||
"key_alias": key_alias,
|
||||
"allowed_routes": BATCH_ROUTES,
|
||||
},
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.master_api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=30,
|
||||
)
|
||||
assert key_response.status_code == 200, (
|
||||
f"Failed to create additional key: {key_response.status_code} - {key_response.text}"
|
||||
)
|
||||
api_key = key_response.json().get("key")
|
||||
print(f"Created additional key {api_key[:20]}... for user {user_id}")
|
||||
return api_key, self.create_openai_client(api_key)
|
||||
|
||||
@pytest.mark.parametrize("model_name", MODEL_NAMES, ids=MODEL_IDS)
|
||||
def test_user_b_cannot_retrieve_user_a_batch(self, tmp_path, model_name):
|
||||
user_a_id, user_a_key, client_A = self.create_user_key_and_client("a")
|
||||
user_b_id, user_b_key, client_B = self.create_user_key_and_client("b")
|
||||
|
||||
# User A creates a batch input file and batch
|
||||
request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
|
||||
batch_input_file = self.create_batch_input_file(client_A, request_file)
|
||||
batch = self.create_batch(client_A, batch_input_file.id)
|
||||
|
||||
# User A retrieves their own batch
|
||||
batch_a = client_A.batches.retrieve(batch_id=batch.id)
|
||||
assert batch_a.id == batch.id, (
|
||||
"User A should be able to retrieve their own batch"
|
||||
)
|
||||
|
||||
# User B cannot retrieve User A's batch
|
||||
try:
|
||||
client_B.batches.retrieve(batch_id=batch.id)
|
||||
pytest.fail("User B should NOT be able to retrieve User A's batch")
|
||||
except openai.PermissionDeniedError as e:
|
||||
assert e.status_code == 403
|
||||
|
||||
@pytest.mark.parametrize("model_name", MODEL_NAMES, ids=MODEL_IDS)
|
||||
def test_user_b_cannot_cancel_user_a_batch(self, tmp_path, model_name):
|
||||
user_a_id, user_a_key, client_A = self.create_user_key_and_client("a")
|
||||
user_b_id, user_b_key, client_B = self.create_user_key_and_client("b")
|
||||
|
||||
# User A creates a batch input file and batch
|
||||
request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
|
||||
batch_input_file = self.create_batch_input_file(client_A, request_file)
|
||||
batch = self.create_batch(client_A, batch_input_file.id)
|
||||
|
||||
# User B cannot cancel User A's batch
|
||||
try:
|
||||
client_B.batches.cancel(batch_id=batch.id)
|
||||
pytest.fail("User B should NOT be able to cancel User A's batch")
|
||||
except openai.PermissionDeniedError as e:
|
||||
assert e.status_code == 403
|
||||
|
||||
@pytest.mark.parametrize("model_name", MODEL_NAMES, ids=MODEL_IDS)
|
||||
def test_user_b_cannot_retrieve_user_a_batch_input_file(self, tmp_path, model_name):
|
||||
user_a_id, user_a_key, client_A = self.create_user_key_and_client("a")
|
||||
user_b_id, user_b_key, client_B = self.create_user_key_and_client("b")
|
||||
|
||||
# User A creates a batch input file and batch
|
||||
request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
|
||||
batch_input_file = self.create_batch_input_file(client_A, request_file)
|
||||
|
||||
# User A retrieves their own file
|
||||
file_a = client_A.files.retrieve(file_id=batch_input_file.id)
|
||||
assert file_a.id == batch_input_file.id, (
|
||||
"User A should be able to retrieve their own file"
|
||||
)
|
||||
|
||||
# User B cannot retrieve User A's file
|
||||
try:
|
||||
client_B.files.retrieve(file_id=batch_input_file.id)
|
||||
pytest.fail("User B should NOT be able to retrieve User A's file")
|
||||
except openai.PermissionDeniedError as e:
|
||||
assert e.status_code == 403
|
||||
|
||||
@pytest.mark.parametrize("model_name", MODEL_NAMES, ids=MODEL_IDS)
|
||||
def test_user_b_cannot_download_user_a_batch_input_file_content(
|
||||
self,
|
||||
tmp_path,
|
||||
model_name,
|
||||
):
|
||||
user_a_id, user_a_key, client_A = self.create_user_key_and_client("a")
|
||||
user_b_id, user_b_key, client_B = self.create_user_key_and_client("b")
|
||||
|
||||
# User A creates a batch input file and batch
|
||||
request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
|
||||
batch_input_file = self.create_batch_input_file(client_A, request_file)
|
||||
|
||||
# User A can download their own file content
|
||||
content_a = client_A.files.content(file_id=batch_input_file.id)
|
||||
assert content_a.text, (
|
||||
"User A should be able to download their own file content"
|
||||
)
|
||||
|
||||
# User B cannot download User A's file content
|
||||
try:
|
||||
client_B.files.content(file_id=batch_input_file.id)
|
||||
pytest.fail("User B should NOT be able to download User A's file content")
|
||||
except openai.PermissionDeniedError as e:
|
||||
assert e.status_code == 403
|
||||
|
||||
@pytest.mark.parametrize("model_name", MODEL_NAMES, ids=MODEL_IDS)
|
||||
def test_user_b_cannot_delete_user_a_batch_input_file(self, tmp_path, model_name):
|
||||
user_a_id, user_a_key, client_A = self.create_user_key_and_client("a")
|
||||
user_b_id, user_b_key, client_B = self.create_user_key_and_client("b")
|
||||
|
||||
# User A creates a batch input file
|
||||
request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
|
||||
batch_input_file = self.create_batch_input_file(client_A, request_file)
|
||||
|
||||
# User B cannot delete User A's file
|
||||
try:
|
||||
client_B.files.delete(file_id=batch_input_file.id)
|
||||
pytest.fail("User B should NOT be able to delete User A's file")
|
||||
except openai.PermissionDeniedError as e:
|
||||
assert e.status_code == 403
|
||||
|
||||
# User A can still retrieve their own file
|
||||
file_a = client_A.files.retrieve(file_id=batch_input_file.id)
|
||||
assert file_a.id == batch_input_file.id, "File should still exist"
|
||||
|
||||
# User A can delete their own file
|
||||
try:
|
||||
client_A.files.delete(file_id=batch_input_file.id)
|
||||
except openai.APIError as e:
|
||||
pytest.fail(f"User A should be able to delete their own file: {e}")
|
||||
|
||||
@pytest.mark.parametrize("model_name", MODEL_NAMES, ids=MODEL_IDS)
|
||||
def test_user_b_cannot_retrieve_user_a_batch_output_file(
|
||||
self,
|
||||
tmp_path,
|
||||
model_name,
|
||||
):
|
||||
user_a_id, user_a_key, client_A = self.create_user_key_and_client("a")
|
||||
user_b_id, user_b_key, client_B = self.create_user_key_and_client("b")
|
||||
|
||||
# User A creates a batch input file and batch
|
||||
request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
|
||||
batch_input_file = self.create_batch_input_file(client_A, request_file)
|
||||
batch = self.create_batch(client_A, batch_input_file.id)
|
||||
|
||||
# Wait for batch to complete
|
||||
completed_batch = self.wait_for_batch_completed(client_A, batch.id)
|
||||
assert completed_batch.output_file_id, "Batch should have an output file"
|
||||
|
||||
# User A retrieves their own output file
|
||||
file_a = client_A.files.retrieve(file_id=completed_batch.output_file_id)
|
||||
assert file_a.id == completed_batch.output_file_id, (
|
||||
"User A should be able to retrieve their own output file"
|
||||
)
|
||||
|
||||
# User B cannot retrieve User A's output file
|
||||
try:
|
||||
client_B.files.retrieve(file_id=completed_batch.output_file_id)
|
||||
pytest.fail("User B should NOT be able to retrieve User A's output file")
|
||||
except openai.PermissionDeniedError as e:
|
||||
assert e.status_code == 403
|
||||
|
||||
@pytest.mark.parametrize("model_name", MODEL_NAMES, ids=MODEL_IDS)
|
||||
def test_user_b_cannot_download_user_a_batch_output_file_content(
|
||||
self,
|
||||
tmp_path,
|
||||
model_name,
|
||||
):
|
||||
user_a_id, user_a_key, client_A = self.create_user_key_and_client("a")
|
||||
user_b_id, user_b_key, client_B = self.create_user_key_and_client("b")
|
||||
|
||||
# User A creates a batch input file and batch
|
||||
request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
|
||||
batch_input_file = self.create_batch_input_file(client_A, request_file)
|
||||
batch = self.create_batch(client_A, batch_input_file.id)
|
||||
|
||||
# Wait for batch to complete
|
||||
completed_batch = self.wait_for_batch_completed(client_A, batch.id)
|
||||
assert completed_batch.output_file_id, "Batch should have an output file"
|
||||
|
||||
# User A can download their own output file content
|
||||
content_a = client_A.files.content(file_id=completed_batch.output_file_id)
|
||||
assert content_a.text, (
|
||||
"User A should be able to download their own output file content"
|
||||
)
|
||||
|
||||
# User B cannot download User A's output file content
|
||||
try:
|
||||
client_B.files.content(file_id=completed_batch.output_file_id)
|
||||
pytest.fail(
|
||||
"User B should NOT be able to download User A's output file content",
|
||||
)
|
||||
except openai.PermissionDeniedError as e:
|
||||
assert e.status_code == 403
|
||||
|
||||
@pytest.mark.parametrize("model_name", MODEL_NAMES, ids=MODEL_IDS)
|
||||
def test_user_b_cannot_delete_user_a_batch_output_file(self, tmp_path, model_name):
|
||||
user_a_id, user_a_key, client_A = self.create_user_key_and_client("a")
|
||||
user_b_id, user_b_key, client_B = self.create_user_key_and_client("b")
|
||||
|
||||
# User A creates a batch input file and batch
|
||||
request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
|
||||
batch_input_file = self.create_batch_input_file(client_A, request_file)
|
||||
batch = self.create_batch(client_A, batch_input_file.id)
|
||||
|
||||
# Wait for batch to complete
|
||||
completed_batch = self.wait_for_batch_completed(client_A, batch.id)
|
||||
assert completed_batch.output_file_id, "Batch should have an output file"
|
||||
|
||||
# User B cannot delete User A's output file
|
||||
try:
|
||||
client_B.files.delete(file_id=completed_batch.output_file_id)
|
||||
pytest.fail("User B should NOT be able to delete User A's output file")
|
||||
except openai.PermissionDeniedError as e:
|
||||
assert e.status_code == 403
|
||||
|
||||
# User A can still retrieve their own output file
|
||||
file_a = client_A.files.retrieve(file_id=completed_batch.output_file_id)
|
||||
assert file_a.id == completed_batch.output_file_id, (
|
||||
"Output file should still exist"
|
||||
)
|
||||
|
||||
# User A can delete their own output file
|
||||
try:
|
||||
client_A.files.delete(file_id=completed_batch.output_file_id)
|
||||
except openai.APIError as e:
|
||||
pytest.fail(f"User A should be able to delete their own output file: {e}")
|
||||
|
||||
@pytest.mark.parametrize("model_name", MODEL_NAMES, ids=MODEL_IDS)
|
||||
def test_user_b_cannot_list_user_a_batches(self, tmp_path, model_name):
|
||||
user_a_id, user_a_key, client_A = self.create_user_key_and_client("a")
|
||||
user_b_id, user_b_key, client_B = self.create_user_key_and_client("b")
|
||||
|
||||
# User A creates a batch input file and batch
|
||||
request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
|
||||
batch_input_file = self.create_batch_input_file(client_A, request_file)
|
||||
batch = self.create_batch(client_A, batch_input_file.id)
|
||||
|
||||
# User A can see their own batch in the list
|
||||
batches_a = client_A.batches.list(
|
||||
limit=10,
|
||||
extra_query={"target_model_names": model_name},
|
||||
)
|
||||
batch_ids_a = [b.id for b in batches_a.data]
|
||||
assert batch.id in batch_ids_a, "User A should see their own batch in the list"
|
||||
|
||||
# User B's batch list should NOT contain User A's batch
|
||||
batches_b = client_B.batches.list(
|
||||
limit=10,
|
||||
extra_query={"target_model_names": model_name},
|
||||
)
|
||||
batch_ids_b = [b.id for b in batches_b.data]
|
||||
assert batch.id not in batch_ids_b, (
|
||||
"User B should NOT see User A's batch in the list"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("model_name", MODEL_NAMES, ids=MODEL_IDS)
|
||||
def test_user_api_keys_are_interchangeable(self, tmp_path, model_name):
|
||||
# Create user with 3 keys
|
||||
user_id, key1, client_Key1 = self.create_user_key_and_client("a")
|
||||
key2, client_Key2 = self.create_key_and_client(user_id, "a2")
|
||||
key3, client_Key3 = self.create_key_and_client(user_id, "a3")
|
||||
|
||||
# Key1: Create batch input file and batch
|
||||
request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
|
||||
batch_input_file = self.create_batch_input_file(client_Key1, request_file)
|
||||
batch = self.create_batch(client_Key1, batch_input_file.id)
|
||||
|
||||
# Key1: Retrieve batch
|
||||
batch_retrieved = client_Key1.batches.retrieve(batch_id=batch.id)
|
||||
assert batch_retrieved.id == batch.id, "Key1 should retrieve its own batch"
|
||||
|
||||
# Key2: Wait for batch completion and retrieve output
|
||||
completed_batch = self.wait_for_batch_completed(client_Key2, batch.id)
|
||||
assert completed_batch.output_file_id, "Batch should have an output file"
|
||||
|
||||
# Key2: Retrieve output file metadata
|
||||
output_file = client_Key2.files.retrieve(file_id=completed_batch.output_file_id)
|
||||
assert output_file.id == completed_batch.output_file_id, (
|
||||
"Key2 should retrieve output file"
|
||||
)
|
||||
|
||||
# Key2: Download output file content
|
||||
output_content = client_Key2.files.content(
|
||||
file_id=completed_batch.output_file_id,
|
||||
)
|
||||
assert output_content.text, "Key2 should download output file content"
|
||||
|
||||
# Key3: List batches and verify batch is visible
|
||||
batches = client_Key3.batches.list(
|
||||
limit=10,
|
||||
extra_query={"target_model_names": model_name},
|
||||
)
|
||||
batch_ids = [b.id for b in batches.data]
|
||||
assert batch.id in batch_ids, "Key3 should see batch in list"
|
||||
|
||||
# Key3: Delete input file
|
||||
try:
|
||||
client_Key3.files.delete(file_id=batch_input_file.id)
|
||||
except openai.APIError as e:
|
||||
pytest.fail(f"Key3 should delete input file: {e}")
|
||||
|
||||
# Key3: Delete output file
|
||||
try:
|
||||
client_Key3.files.delete(file_id=completed_batch.output_file_id)
|
||||
except openai.APIError as e:
|
||||
pytest.fail(f"Key3 should delete output file: {e}")
|
||||
|
||||
# Key1: Create another batch for cancellation test
|
||||
request_file2 = self.create_batch_request_file_on_disk(tmp_path, model_name)
|
||||
batch_input_file2 = self.create_batch_input_file(client_Key1, request_file2)
|
||||
batch2 = self.create_batch(client_Key1, batch_input_file2.id)
|
||||
|
||||
# Key3: Cancel the batch created by Key1
|
||||
try:
|
||||
cancelled_batch = client_Key3.batches.cancel(batch_id=batch2.id)
|
||||
assert cancelled_batch.id == batch2.id, (
|
||||
"Key3 should cancel batch created by Key1"
|
||||
)
|
||||
except openai.BadRequestError:
|
||||
pass # Batch may have already completed
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user