Spaces:
Paused
Paused
| import pytest | |
| import httpx | |
| from openai import OpenAI | |
| import os | |
| import time | |
| import random | |
| # Configure the OpenAI client to point to the local server. | |
| # The API key is required by the library but not used by the local server. | |
| client = OpenAI( | |
| base_url="http://localhost:8000", | |
| api_key=os.environ.get("OPENAI_API_KEY", "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx") | |
| ) | |
| async def test_list_models_openai(): | |
| """Test the /models endpoint to get a list of available models.""" | |
| print("\nTesting /models endpoint...") | |
| try: | |
| start_time = time.time() | |
| models_response = client.models.list() | |
| end_time = time.time() | |
| print(f"Received response from /models in {end_time - start_time:.2f}s") | |
| assert models_response is not None | |
| assert len(models_response.data) > 0, "No models were returned from the /models endpoint." | |
| print(f"Found {len(models_response.data)} models:") | |
| for model in models_response.data: | |
| print(f" - Model ID: {model.id}") | |
| assert isinstance(model.id, str) | |
| assert model.object == "model" | |
| except Exception as e: | |
| pytest.fail(f"/models endpoint test failed using openai lib: {e}") | |
| async def test_chat_completion_streaming_openai(): | |
| """ | |
| Test the /chat/completions endpoint for streaming requests. | |
| This test first fetches an available model from the /models endpoint. | |
| """ | |
| # Step 1: Get the list of available models first to make the test robust. | |
| print("\nFetching available models before chat completion test...") | |
| try: | |
| models_response = client.models.list() | |
| assert len(models_response.data) > 0, "Cannot run chat completion test: no models available." | |
| # Use the first model from the list for the test. | |
| model_to_test = models_response.data[random.randint(0, len(models_response.data) - 1)].id | |
| print(f"Will use model '{model_to_test}' for the test.") | |
| except Exception as e: | |
| pytest.fail(f"Failed to fetch models before running chat completion test: {e}") | |
| # Step 2: Use the fetched model to run the chat completion test. | |
| messages = [{"role": "user", "content": "squirrel pet simulator in python with emojis"}] | |
| print(f"\nTesting streaming chat completion with model '{model_to_test}'...") | |
| accumulated_content = "" | |
| deltas_received = 0 | |
| start_time = time.time() | |
| try: | |
| stream = client.chat.completions.create( | |
| model=model_to_test, | |
| messages=messages, | |
| stream=True | |
| ) | |
| for chunk in stream: | |
| assert chunk.id is not None | |
| assert chunk.object == "chat.completion.chunk" | |
| assert len(chunk.choices) > 0 | |
| choice = chunk.choices[0] | |
| delta = choice.delta | |
| if delta.role is not None: | |
| print(f"\nReceived role: {delta.role}") | |
| assert delta.role == "assistant" | |
| if delta.content is not None: | |
| delta_content = delta.content | |
| assert isinstance(delta_content, str) | |
| accumulated_content += delta_content | |
| deltas_received += 1 | |
| print(delta_content, end="", flush=True) | |
| if choice.finish_reason is not None: | |
| print(f"\nReceived finish_reason: {choice.finish_reason}") | |
| assert choice.finish_reason == "stop" | |
| break | |
| end_time = time.time() | |
| print(f"\n--- Streaming complete ---") | |
| print(f"Total deltas received: {deltas_received}") | |
| print(f"Full response ({len(accumulated_content)} chars) received in {end_time - start_time:.2f}s:") | |
| print(accumulated_content) | |
| assert deltas_received > 0, "No content deltas were received." | |
| assert len(accumulated_content) > 0, "Accumulated content is empty." | |
| assert "squirrel" in accumulated_content.lower(), "Story does not seem to be about a squirrel." | |
| except Exception as e: | |
| pytest.fail(f"Streaming chat completion failed: {e}") | |