mirror of
https://github.com/steelcityamir/safe-content-ai.git
synced 2024-10-03 21:38:34 +03:00
4ce040a5b2
* Update main.py * Update ci.yml * Update README.md * Update requirements.txt * Create models.py * Update main.py * Update test_main.py * Update README.md * Update models.py * Update test_main.py * Update main.py * Update main.py * Update models.py * Update main.py
84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
"""API Tests"""
|
|
|
|
import hashlib
|
|
from unittest.mock import patch
|
|
from fastapi.testclient import TestClient
|
|
from main import app, cache
|
|
|
|
client = TestClient(app)
|
|
|
|
FILE_NAME = "sunflower.jpg"
|
|
IMAGE_HASH = "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824"
|
|
|
|
|
|
def compute_file_hash(file_path):
|
|
"""Utility function to compute SHA-256 hash of a file"""
|
|
hash_sha256 = hashlib.sha256()
|
|
with open(file_path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(4096), b""):
|
|
hash_sha256.update(chunk)
|
|
return hash_sha256.hexdigest()
|
|
|
|
|
|
def test_read_main():
|
|
"""Tests that POST /v1/detect returns 200 OK with valid request body"""
|
|
with open(FILE_NAME, "rb") as file:
|
|
response = client.post(
|
|
"/v1/detect",
|
|
files={"file": (FILE_NAME, file, "image/jpeg")},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json() == {
|
|
"file_name": FILE_NAME,
|
|
"is_nsfw": False,
|
|
"confidence_percentage": 100.0,
|
|
}
|
|
|
|
|
|
def test_invalid_input():
|
|
"""Tests that POST /v1/detect returns 400 with empty request body"""
|
|
response = client.post("/v1/detect", files={})
|
|
assert response.status_code == 400
|
|
|
|
|
|
def test_cache_hit():
|
|
"""Tests that the endpoint returns a cached response when an image hash matches"""
|
|
# Compute the hash of the test file
|
|
image_hash = compute_file_hash(FILE_NAME)
|
|
cached_response = {
|
|
"file_name": FILE_NAME,
|
|
"is_nsfw": False,
|
|
"confidence_percentage": 100.0,
|
|
}
|
|
|
|
with patch.dict(cache, {image_hash: cached_response}), patch(
|
|
"main.logging.info"
|
|
) as mock_logging:
|
|
# Post the file to the endpoint as would happen in actual use
|
|
with open(FILE_NAME, "rb") as file:
|
|
response = client.post(
|
|
"/v1/detect",
|
|
files={"file": (FILE_NAME, file, "image/jpeg")},
|
|
)
|
|
|
|
# Verify that the response is from the cache
|
|
assert response.status_code == 200
|
|
assert response.json() == cached_response
|
|
|
|
# Ensure logging was called correctly
|
|
mock_logging.assert_called_with("Returning cached entry for %s", FILE_NAME)
|
|
|
|
def test_detect_urls():
|
|
"""Tests that POST /v1/detect/urls returns 200 OK with valid request body"""
|
|
urls = [
|
|
"https://raw.githubusercontent.com/steelcityamir/safe-content-ai/main/sunflower.jpg",
|
|
]
|
|
response = client.post("/v1/detect/urls", json={"urls": urls})
|
|
assert response.status_code == 200
|
|
assert len(response.json()) == len(urls)
|
|
assert response.json()[0] == {
|
|
"url": urls[0],
|
|
"is_nsfw": False,
|
|
"confidence_percentage": 100.0,
|
|
}
|