84c09e41ef
Add 503/403 auth tests for the /train-sender endpoint, matching the pattern already used for /train and /segtrain. Also surface test_sender_registry.py in CI (it needs no ML stack) and add pytest-asyncio to the install step. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
109 lines
4.3 KiB
Python
109 lines
4.3 KiB
Python
"""Tests for /train and /segtrain endpoint authentication."""
|
|
|
|
import io
|
|
import zipfile
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
from httpx import ASGITransport, AsyncClient
|
|
|
|
from main import app
|
|
|
|
|
|
def _minimal_zip() -> bytes:
|
|
"""Return a ZIP with one .xml file so endpoint validation passes."""
|
|
buf = io.BytesIO()
|
|
with zipfile.ZipFile(buf, "w") as zf:
|
|
zf.writestr("page_01.xml", "<PcGts/>")
|
|
return buf.getvalue()
|
|
|
|
|
|
# ─── Missing TRAINING_TOKEN → fail closed ─────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_train_returns_503_when_training_token_not_configured():
|
|
"""POST /train must return 503 when TRAINING_TOKEN env var is empty.
|
|
|
|
An empty token means the service was started without training configured.
|
|
Allowing requests through would grant unauthenticated access to the
|
|
training endpoint, contradicting the principle of failing closed.
|
|
"""
|
|
with patch("main.TRAINING_TOKEN", ""), \
|
|
patch("main._models_ready", True):
|
|
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
|
response = await client.post(
|
|
"/train",
|
|
files={"file": ("training.zip", _minimal_zip(), "application/zip")},
|
|
)
|
|
|
|
assert response.status_code == 503
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_segtrain_returns_503_when_training_token_not_configured():
|
|
"""POST /segtrain must return 503 when TRAINING_TOKEN env var is empty."""
|
|
with patch("main.TRAINING_TOKEN", ""), \
|
|
patch("main._models_ready", True):
|
|
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
|
response = await client.post(
|
|
"/segtrain",
|
|
files={"file": ("training.zip", _minimal_zip(), "application/zip")},
|
|
)
|
|
|
|
assert response.status_code == 503
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_train_returns_403_when_token_wrong():
|
|
"""POST /train must return 403 when TRAINING_TOKEN is set but header is wrong."""
|
|
with patch("main.TRAINING_TOKEN", "secret-token"), \
|
|
patch("main._models_ready", True):
|
|
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
|
response = await client.post(
|
|
"/train",
|
|
files={"file": ("training.zip", _minimal_zip(), "application/zip")},
|
|
headers={"X-Training-Token": "wrong-token"},
|
|
)
|
|
|
|
assert response.status_code == 403
|
|
|
|
|
|
# ─── /train-sender authentication ─────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_train_sender_returns_503_when_training_token_not_configured():
|
|
"""POST /train-sender must return 503 when TRAINING_TOKEN env var is empty.
|
|
|
|
An empty token means the service was started without training configured.
|
|
Allowing requests through would grant unauthenticated access to the
|
|
training endpoint, contradicting the principle of failing closed.
|
|
"""
|
|
with patch("main.TRAINING_TOKEN", ""), \
|
|
patch("main._models_ready", True):
|
|
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
|
response = await client.post(
|
|
"/train-sender",
|
|
data={"output_model_path": "/app/models/sender_test.mlmodel"},
|
|
files={"file": ("training.zip", _minimal_zip(), "application/zip")},
|
|
)
|
|
|
|
assert response.status_code == 503
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_train_sender_returns_403_when_token_wrong():
|
|
"""POST /train-sender must return 403 when TRAINING_TOKEN is set but header is wrong."""
|
|
with patch("main.TRAINING_TOKEN", "secret-token"), \
|
|
patch("main._models_ready", True):
|
|
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
|
response = await client.post(
|
|
"/train-sender",
|
|
data={"output_model_path": "/app/models/sender_test.mlmodel"},
|
|
files={"file": ("training.zip", _minimal_zip(), "application/zip")},
|
|
headers={"X-Training-Token": "wrong-token"},
|
|
)
|
|
|
|
assert response.status_code == 403
|