mirror of
https://github.com/azaion/detections.git
synced 2026-04-22 10:46:31 +00:00
Oleksandr Bezdieniezhnykh
097811a67b
- Pin all deps; h11==0.16.0 (CVE-2025-43859), python-multipart>=1.3.1 (CVE-2026-28356), PyJWT==2.12.1
- Add HMAC JWT verification (require_auth FastAPI dependency, JWT_SECRET-gated)
- Fix TokenManager._refresh() to use ADMIN_API_URL instead of ANNOTATIONS_URL
- Rename POST /detect → POST /detect/image (image-only, rejects video files)
- Replace global SSE stream with per-job SSE: GET /detect/{media_id} with event replay buffer
- Apply require_auth to all 4 protected endpoints
- Fix on_annotation/on_status closure to use mutable current_id for correct post-upload event routing
- Add non-root appuser to Dockerfile and Dockerfile.gpu
- Add JWT_SECRET to e2e/docker-compose.test.yml and run-tests.sh
- Update all e2e tests and unit tests for new endpoints and HMAC token signing
- 64/64 tests pass
Made-with: Cursor
442 lines
12 KiB
Python
442 lines
12 KiB
Python
import asyncio
|
|
import base64
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
class TestStreamingBuffer:
|
|
def test_sequential_write_read(self):
|
|
# Arrange
|
|
from streaming_buffer import StreamingBuffer
|
|
|
|
buf = StreamingBuffer()
|
|
try:
|
|
buf.append(b"hello")
|
|
buf.append(b" world")
|
|
buf.close_writer()
|
|
# Act
|
|
result = buf.read(-1)
|
|
# Assert
|
|
assert result == b"hello world"
|
|
finally:
|
|
buf.close()
|
|
os.unlink(buf.path)
|
|
|
|
def test_read_blocks_until_data_available(self):
|
|
# Arrange
|
|
from streaming_buffer import StreamingBuffer
|
|
|
|
buf = StreamingBuffer()
|
|
results = []
|
|
|
|
def writer():
|
|
time.sleep(0.1)
|
|
buf.append(b"data")
|
|
buf.close_writer()
|
|
|
|
t = threading.Thread(target=writer)
|
|
t.start()
|
|
|
|
# Act
|
|
results.append(buf.read(4))
|
|
t.join(timeout=5)
|
|
|
|
# Assert
|
|
assert results == [b"data"]
|
|
buf.close()
|
|
os.unlink(buf.path)
|
|
|
|
def test_read_returns_empty_on_eof(self):
|
|
# Arrange
|
|
from streaming_buffer import StreamingBuffer
|
|
|
|
buf = StreamingBuffer()
|
|
buf.close_writer()
|
|
|
|
# Act
|
|
result = buf.read(1024)
|
|
|
|
# Assert
|
|
assert result == b""
|
|
buf.close()
|
|
os.unlink(buf.path)
|
|
|
|
def test_concurrent_write_read_chunked(self):
|
|
# Arrange
|
|
from streaming_buffer import StreamingBuffer
|
|
|
|
buf = StreamingBuffer()
|
|
chunks_written = [b"aaa", b"bbb", b"ccc"]
|
|
read_data = bytearray()
|
|
|
|
def writer():
|
|
for c in chunks_written:
|
|
time.sleep(0.02)
|
|
buf.append(c)
|
|
buf.close_writer()
|
|
|
|
def reader():
|
|
while True:
|
|
chunk = buf.read(1024)
|
|
if not chunk:
|
|
break
|
|
read_data.extend(chunk)
|
|
|
|
wt = threading.Thread(target=writer)
|
|
rt = threading.Thread(target=reader)
|
|
|
|
# Act
|
|
wt.start()
|
|
rt.start()
|
|
wt.join(timeout=5)
|
|
rt.join(timeout=5)
|
|
|
|
# Assert
|
|
assert bytes(read_data) == b"aaabbbccc"
|
|
buf.close()
|
|
os.unlink(buf.path)
|
|
|
|
def test_seek_set_and_reread(self):
|
|
# Arrange
|
|
from streaming_buffer import StreamingBuffer
|
|
|
|
buf = StreamingBuffer()
|
|
buf.append(b"0123456789")
|
|
buf.close_writer()
|
|
|
|
# Act
|
|
buf.read(5)
|
|
buf.seek(2, 0)
|
|
result = buf.read(3)
|
|
|
|
# Assert
|
|
assert result == b"234"
|
|
buf.close()
|
|
os.unlink(buf.path)
|
|
|
|
def test_seek_end_blocks_until_eof(self):
|
|
# Arrange
|
|
from streaming_buffer import StreamingBuffer
|
|
|
|
buf = StreamingBuffer()
|
|
positions = []
|
|
|
|
def writer():
|
|
time.sleep(0.1)
|
|
buf.append(b"abcdef")
|
|
buf.close_writer()
|
|
|
|
t = threading.Thread(target=writer)
|
|
t.start()
|
|
|
|
# Act
|
|
pos = buf.seek(0, 2)
|
|
positions.append(pos)
|
|
t.join(timeout=5)
|
|
|
|
# Assert
|
|
assert positions[0] == 6
|
|
buf.close()
|
|
os.unlink(buf.path)
|
|
|
|
def test_tell_tracks_position(self):
|
|
# Arrange
|
|
from streaming_buffer import StreamingBuffer
|
|
|
|
buf = StreamingBuffer()
|
|
buf.append(b"data")
|
|
buf.close_writer()
|
|
|
|
# Assert
|
|
assert buf.tell() == 0
|
|
buf.read(2)
|
|
assert buf.tell() == 2
|
|
buf.close()
|
|
os.unlink(buf.path)
|
|
|
|
def test_file_persisted_to_disk(self):
|
|
# Arrange
|
|
from streaming_buffer import StreamingBuffer
|
|
|
|
buf = StreamingBuffer()
|
|
payload = b"x" * 10000
|
|
|
|
# Act
|
|
buf.append(payload)
|
|
buf.close_writer()
|
|
|
|
# Assert
|
|
with open(buf.path, "rb") as f:
|
|
assert f.read() == payload
|
|
buf.close()
|
|
os.unlink(buf.path)
|
|
|
|
def test_written_property(self):
|
|
# Arrange
|
|
from streaming_buffer import StreamingBuffer
|
|
|
|
buf = StreamingBuffer()
|
|
buf.append(b"abc")
|
|
buf.append(b"defgh")
|
|
buf.close_writer()
|
|
|
|
# Assert
|
|
assert buf.written == 8
|
|
buf.close()
|
|
os.unlink(buf.path)
|
|
|
|
def test_seekable_readable(self):
|
|
# Arrange
|
|
from streaming_buffer import StreamingBuffer
|
|
|
|
buf = StreamingBuffer()
|
|
buf.close_writer()
|
|
|
|
# Assert
|
|
assert buf.seekable() is True
|
|
assert buf.readable() is True
|
|
assert buf.writable() is False
|
|
buf.close()
|
|
os.unlink(buf.path)
|
|
|
|
|
|
class TestMediaContentHashFromFile:
|
|
def test_small_file_matches_bytes_version(self):
|
|
# Arrange
|
|
from media_hash import compute_media_content_hash, compute_media_content_hash_from_file
|
|
|
|
data = b"hello world"
|
|
with tempfile.NamedTemporaryFile(delete=False) as f:
|
|
f.write(data)
|
|
path = f.name
|
|
|
|
# Act
|
|
hash_bytes = compute_media_content_hash(data)
|
|
hash_file = compute_media_content_hash_from_file(path)
|
|
|
|
# Assert
|
|
assert hash_file == hash_bytes
|
|
os.unlink(path)
|
|
|
|
def test_large_file_matches_bytes_version(self):
|
|
# Arrange
|
|
from media_hash import compute_media_content_hash, compute_media_content_hash_from_file
|
|
|
|
data = os.urandom(50_000)
|
|
with tempfile.NamedTemporaryFile(delete=False) as f:
|
|
f.write(data)
|
|
path = f.name
|
|
|
|
# Act
|
|
hash_bytes = compute_media_content_hash(data)
|
|
hash_file = compute_media_content_hash_from_file(path)
|
|
|
|
# Assert
|
|
assert hash_file == hash_bytes
|
|
os.unlink(path)
|
|
|
|
def test_virtual_flag(self):
|
|
# Arrange
|
|
from media_hash import compute_media_content_hash_from_file
|
|
|
|
with tempfile.NamedTemporaryFile(delete=False) as f:
|
|
f.write(b"test")
|
|
path = f.name
|
|
|
|
# Act
|
|
normal = compute_media_content_hash_from_file(path, virtual=False)
|
|
virtual = compute_media_content_hash_from_file(path, virtual=True)
|
|
|
|
# Assert
|
|
assert virtual == f"V{normal}"
|
|
os.unlink(path)
|
|
|
|
def test_exact_boundary_3072_bytes(self):
|
|
# Arrange
|
|
from media_hash import compute_media_content_hash, compute_media_content_hash_from_file
|
|
|
|
data = os.urandom(3072)
|
|
with tempfile.NamedTemporaryFile(delete=False) as f:
|
|
f.write(data)
|
|
path = f.name
|
|
|
|
# Act
|
|
hash_bytes = compute_media_content_hash(data)
|
|
hash_file = compute_media_content_hash_from_file(path)
|
|
|
|
# Assert
|
|
assert hash_file == hash_bytes
|
|
os.unlink(path)
|
|
|
|
|
|
def _access_jwt(sub: str = "u1") -> str:
|
|
import jwt as pyjwt
|
|
secret = os.environ.get("JWT_SECRET", "")
|
|
if secret:
|
|
return pyjwt.encode(
|
|
{"exp": int(time.time()) + 3600, "sub": sub},
|
|
secret,
|
|
algorithm="HS256",
|
|
)
|
|
raw = json.dumps(
|
|
{"exp": int(time.time()) + 3600, "sub": sub}, separators=(",", ":")
|
|
).encode()
|
|
payload = base64.urlsafe_b64encode(raw).decode().rstrip("=")
|
|
return f"h.{payload}.s"
|
|
|
|
|
|
class _FakeInfStream:
|
|
is_engine_ready = True
|
|
|
|
def run_detect_video_stream(
|
|
self, readable, ai_cfg, media_name, on_annotation, status_callback=None
|
|
):
|
|
while True:
|
|
chunk = readable.read(4096)
|
|
if not chunk:
|
|
break
|
|
if status_callback:
|
|
status_callback(media_name, 0)
|
|
|
|
def run_detect_video(self, *a, **kw):
|
|
pass
|
|
|
|
def run_detect_image(self, *a, **kw):
|
|
pass
|
|
|
|
|
|
class TestDetectVideoEndpoint:
|
|
@pytest.fixture(autouse=True)
|
|
def reset_inference(self):
|
|
import main
|
|
main.inference = None
|
|
yield
|
|
main.inference = None
|
|
|
|
def test_streaming_upload_returns_started(self):
|
|
# Arrange
|
|
import main
|
|
from media_hash import compute_media_content_hash
|
|
|
|
video_body = b"fake-video-" * 200
|
|
content_hash = compute_media_content_hash(video_body)
|
|
mock_post = MagicMock()
|
|
mock_post.return_value.status_code = 201
|
|
mock_put = MagicMock()
|
|
mock_put.return_value.status_code = 204
|
|
token = _access_jwt()
|
|
|
|
with tempfile.TemporaryDirectory() as vd:
|
|
os.environ["VIDEOS_DIR"] = vd
|
|
from fastapi.testclient import TestClient
|
|
client = TestClient(main.app)
|
|
with (
|
|
patch.object(main, "get_inference", return_value=_FakeInfStream()),
|
|
patch.object(main.http_requests, "post", mock_post),
|
|
patch.object(main.http_requests, "put", mock_put),
|
|
):
|
|
# Act
|
|
r = client.post(
|
|
"/detect/video",
|
|
content=video_body,
|
|
headers={
|
|
"X-Filename": "test.mp4",
|
|
"Authorization": f"Bearer {token}",
|
|
},
|
|
)
|
|
# Assert
|
|
assert r.status_code == 200
|
|
data = r.json()
|
|
assert data["status"] == "started"
|
|
assert data["mediaId"] == content_hash
|
|
stored = os.path.join(vd, f"{content_hash}.mp4")
|
|
assert os.path.isfile(stored)
|
|
with open(stored, "rb") as f:
|
|
assert f.read() == video_body
|
|
|
|
def test_non_auth_cleanup(self):
|
|
# Arrange
|
|
import main
|
|
|
|
video_body = b"noauth-vid-" * 100
|
|
with tempfile.TemporaryDirectory() as vd:
|
|
os.environ["VIDEOS_DIR"] = vd
|
|
from fastapi.testclient import TestClient
|
|
client = TestClient(main.app)
|
|
with (
|
|
patch.object(main, "JWT_SECRET", ""),
|
|
patch.object(main, "get_inference", return_value=_FakeInfStream()),
|
|
):
|
|
# Act
|
|
r = client.post(
|
|
"/detect/video",
|
|
content=video_body,
|
|
headers={"X-Filename": "test.mp4"},
|
|
)
|
|
# Assert
|
|
assert r.status_code == 200
|
|
assert r.json()["status"] == "started"
|
|
|
|
def test_rejects_non_video_extension(self):
|
|
# Arrange
|
|
import main
|
|
|
|
from fastapi.testclient import TestClient
|
|
client = TestClient(main.app)
|
|
|
|
# Act — patch JWT_SECRET to "" so auth does not block the extension check
|
|
with patch.object(main, "JWT_SECRET", ""):
|
|
r = client.post(
|
|
"/detect/video",
|
|
content=b"data",
|
|
headers={"X-Filename": "photo.jpg"},
|
|
)
|
|
|
|
# Assert
|
|
assert r.status_code == 400
|
|
|
|
def test_data_flows_through_streaming_buffer(self):
|
|
# Arrange
|
|
import main
|
|
from streaming_buffer import StreamingBuffer
|
|
|
|
received_chunks = []
|
|
|
|
class _CaptureInf(_FakeInfStream):
|
|
def run_detect_video_stream(
|
|
self, readable, ai_cfg, media_name, on_annotation, status_callback=None
|
|
):
|
|
while True:
|
|
chunk = readable.read(4096)
|
|
if not chunk:
|
|
break
|
|
received_chunks.append(chunk)
|
|
|
|
video_body = b"A" * 10000
|
|
with tempfile.TemporaryDirectory() as vd:
|
|
os.environ["VIDEOS_DIR"] = vd
|
|
from fastapi.testclient import TestClient
|
|
client = TestClient(main.app)
|
|
token = _access_jwt()
|
|
with patch.object(main, "get_inference", return_value=_CaptureInf()):
|
|
# Act
|
|
r = client.post(
|
|
"/detect/video",
|
|
content=video_body,
|
|
headers={
|
|
"X-Filename": "v.mp4",
|
|
"Authorization": f"Bearer {token}",
|
|
},
|
|
)
|
|
|
|
# Assert
|
|
assert r.status_code == 200
|
|
all_received = b"".join(received_chunks)
|
|
assert all_received == video_body
|