gunicorn/examples/embedding_service/test_embedding.py
Benoit Chesneau 95b7ffeeaa chore: prepare release 25.0.2
- Bump version to 25.0.2
- Update copyright year to 2026 in LICENSE and NOTICE
- Add license headers to all Python source files
- Add changelog entry for 25.0.2
2026-02-06 08:21:18 +01:00

38 lines
1.1 KiB
Python

#
# This file is part of gunicorn released under the MIT license.
# See the NOTICE for more information.
import os
import requests
import numpy as np
def test_embedding_endpoint():
base_url = os.environ.get("EMBEDDING_SERVICE_URL", "http://127.0.0.1:8000")
url = f"{base_url}/embed"
# Test single text
response = requests.post(url, json={"texts": ["Hello world"]})
assert response.status_code == 200
data = response.json()
assert len(data["embeddings"]) == 1
assert len(data["embeddings"][0]) == 384 # MiniLM dimension
# Test batch
texts = ["First sentence", "Second sentence", "Third one"]
response = requests.post(url, json={"texts": texts})
assert response.status_code == 200
data = response.json()
assert len(data["embeddings"]) == 3
# Test similarity (same text = same embedding)
response = requests.post(url, json={"texts": ["test", "test"]})
emb1, emb2 = response.json()["embeddings"]
assert np.allclose(emb1, emb2, rtol=1e-5, atol=1e-6)
print("All tests passed!")
if __name__ == "__main__":
test_embedding_endpoint()