Add E2E tests, fix bugs

Made-with: Cursor

e2e/conftest.py

@@ -0,0 +1,68 @@
+import os
+import subprocess
+import time
+
+import boto3
+import pytest
+import requests
+from botocore.config import Config
+from botocore.exceptions import ClientError
+
+COMPOSE_FILE = os.path.join(os.path.dirname(__file__), "docker-compose.test.yml")
+
+
+@pytest.fixture(scope="session")
+def base_url():
+    return os.environ.get("LOADER_URL", "http://localhost:8080").rstrip("/")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def _reset_loader(base_url):
+    subprocess.run(
+        ["docker", "compose", "-f", COMPOSE_FILE, "restart", "system-under-test"],
+        capture_output=True, timeout=30,
+    )
+
+    endpoint = os.environ.get("MINIO_URL", "http://localhost:9000")
+    s3 = boto3.client(
+        "s3",
+        endpoint_url=endpoint,
+        aws_access_key_id="minioadmin",
+        aws_secret_access_key="minioadmin",
+        config=Config(signature_version="s3v4"),
+        region_name="us-east-1",
+    )
+    for bucket in ["models"]:
+        try:
+            s3.head_bucket(Bucket=bucket)
+            for obj in s3.list_objects_v2(Bucket=bucket).get("Contents", []):
+                s3.delete_object(Bucket=bucket, Key=obj["Key"])
+        except ClientError:
+            s3.create_bucket(Bucket=bucket)
+
+    session = requests.Session()
+    deadline = time.monotonic() + 30
+    while time.monotonic() < deadline:
+        try:
+            if session.get(f"{base_url}/health", timeout=2).status_code == 200:
+                break
+        except Exception:
+            pass
+        time.sleep(1)
+
+
+@pytest.fixture
+def api_client():
+    return requests.Session()
+
+
+@pytest.fixture
+def logged_in_client(base_url, api_client):
+    email = os.environ.get("TEST_EMAIL", "test@azaion.com")
+    password = os.environ.get("TEST_PASSWORD", "testpass")
+    response = api_client.post(
+        f"{base_url}/login",
+        json={"email": email, "password": password},
+    )
+    response.raise_for_status()
+    return api_client

e2e/docker-compose.test.yml

@@ -0,0 +1,43 @@
+services:
+  mock-api:
+    build: ./mocks/mock_api
+    ports:
+      - "9090:9090"
+    environment:
+      MOCK_CDN_HOST: http://mock-cdn:9000
+    networks:
+      - e2e-net
+
+  mock-cdn:
+    image: minio/minio:latest
+    command: server /data --console-address ":9001"
+    environment:
+      MINIO_ROOT_USER: minioadmin
+      MINIO_ROOT_PASSWORD: minioadmin
+    ports:
+      - "9000:9000"
+    networks:
+      - e2e-net
+
+  system-under-test:
+    build:
+      context: ..
+      dockerfile: Dockerfile
+    command: bash -c "rm -rf /app/models/* && mkdir -p /app/models && python -m uvicorn main:app --host 0.0.0.0 --port 8080"
+    ports:
+      - "8080:8080"
+    depends_on:
+      - mock-api
+      - mock-cdn
+    environment:
+      RESOURCE_API_URL: http://mock-api:9090
+      IMAGES_PATH: /tmp/test.enc
+      API_VERSION: test
+    volumes:
+      - /var/run/docker.sock:/var/run/docker.sock
+    networks:
+      - e2e-net
+
+networks:
+  e2e-net:
+    driver: bridge

e2e/mocks/mock_api/Dockerfile

@@ -0,0 +1,7 @@
+FROM python:3.11-slim
+WORKDIR /app
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+COPY app.py .
+EXPOSE 9090
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "9090"]

e2e/mocks/mock_api/app.py

@@ -0,0 +1,119 @@
+import base64
+import hashlib
+import os
+import secrets
+import uuid
+
+import jwt
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.primitives import padding
+from fastapi import FastAPI, File, Request, UploadFile
+from fastapi.responses import JSONResponse, Response
+from pydantic import BaseModel
+
+VALID_EMAIL = os.environ.get("MOCK_VALID_EMAIL", "test@azaion.com")
+VALID_PASSWORD = os.environ.get("MOCK_VALID_PASSWORD", "testpass")
+JWT_SECRET = os.environ.get("MOCK_JWT_SECRET", "e2e-mock-jwt-secret")
+CDN_HOST = os.environ.get("MOCK_CDN_HOST", "http://mock-cdn:9000")
+
+CDN_CONFIG_YAML = (
+    f"host: {CDN_HOST}\n"
+    "downloader_access_key: minioadmin\n"
+    "downloader_access_secret: minioadmin\n"
+    "uploader_access_key: minioadmin\n"
+    "uploader_access_secret: minioadmin\n"
+)
+
+uploaded_files: dict[str, bytes] = {}
+
+app = FastAPI()
+
+
+class LoginBody(BaseModel):
+    email: str
+    password: str
+
+
+def _calc_hash(key: str) -> str:
+    h = hashlib.sha384(key.encode("utf-8")).digest()
+    return base64.b64encode(h).decode("utf-8")
+
+
+def _encrypt(plaintext: bytes, key: str) -> bytes:
+    aes_key = hashlib.sha256(key.encode("utf-8")).digest()
+    iv = os.urandom(16)
+    cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend())
+    encryptor = cipher.encryptor()
+    padder = padding.PKCS7(128).padder()
+    padded = padder.update(plaintext) + padder.finalize()
+    ciphertext = encryptor.update(padded) + encryptor.finalize()
+    return iv + ciphertext
+
+
+@app.post("/login")
+def login(body: LoginBody):
+    if body.email == VALID_EMAIL and body.password == VALID_PASSWORD:
+        token = jwt.encode(
+            {
+                "nameid": str(uuid.uuid4()),
+                "unique_name": body.email,
+                "role": "Admin",
+            },
+            JWT_SECRET,
+            algorithm="HS256",
+        )
+        if isinstance(token, bytes):
+            token = token.decode("ascii")
+        return {"token": token}
+    return JSONResponse(
+        status_code=409,
+        content={"ErrorCode": "AUTH_FAILED", "Message": "Invalid credentials"},
+    )
+
+
+@app.post("/resources/get/{folder:path}")
+async def resources_get(folder: str, request: Request):
+    body = await request.json()
+    hardware = body.get("hardware", "")
+    password = body.get("password", "")
+    filename = body.get("fileName", "")
+
+    hw_hash = _calc_hash(f"Azaion_{hardware}_%$$$)0_")
+    enc_key = _calc_hash(f"{VALID_EMAIL}-{password}-{hw_hash}-#%@AzaionKey@%#---")
+
+    if filename == "cdn.yaml":
+        encrypted = _encrypt(CDN_CONFIG_YAML.encode("utf-8"), enc_key)
+        return Response(content=encrypted, media_type="application/octet-stream")
+
+    storage_key = f"{folder}/{filename}" if folder else filename
+    if storage_key in uploaded_files:
+        encrypted = _encrypt(uploaded_files[storage_key], enc_key)
+        return Response(content=encrypted, media_type="application/octet-stream")
+
+    encrypted = _encrypt(b"\x00" * 32, enc_key)
+    return Response(content=encrypted, media_type="application/octet-stream")
+
+
+@app.post("/resources/{folder}")
+async def resources_upload(folder: str, data: UploadFile = File(...)):
+    content = await data.read()
+    storage_key = f"{folder}/{data.filename}"
+    uploaded_files[storage_key] = content
+    return Response(status_code=200)
+
+
+@app.get("/resources/list/{folder}")
+def resources_list(folder: str, search: str = ""):
+    return []
+
+
+@app.get("/binary-split/key-fragment")
+def binary_split_key_fragment():
+    return Response(content=secrets.token_bytes(16), media_type="application/octet-stream")
+
+
+@app.post("/resources/check")
+async def resources_check(request: Request):
+    await request.body()
+    return Response(status_code=200)

e2e/mocks/mock_api/requirements.txt

@@ -0,0 +1,5 @@
+fastapi
+uvicorn
+pyjwt
+python-multipart
+cryptography

e2e/pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+addopts = -v

e2e/requirements.txt

@@ -0,0 +1,3 @@
+pytest
+requests
+boto3

e2e/tests/test_auth.py

@@ -0,0 +1,59 @@
+def test_status_unauthenticated(base_url, api_client):
+    # Act
+    response = api_client.get(f"{base_url}/status")
+
+    # Assert
+    assert response.status_code == 200
+    assert response.json()["authenticated"] is False
+
+
+def test_download_unauthenticated(base_url, api_client):
+    # Arrange
+    url = f"{base_url}/load/testmodel"
+    body = {"filename": "testmodel", "folder": "models"}
+
+    # Act
+    response = api_client.post(url, json=body)
+
+    # Assert
+    assert response.status_code == 500
+
+
+def test_login_invalid_credentials(base_url, api_client):
+    # Arrange
+    payload = {"email": "wrong@example.com", "password": "wrong"}
+
+    # Act
+    response = api_client.post(f"{base_url}/login", json=payload)
+
+    # Assert
+    assert response.status_code == 401
+
+
+def test_login_empty_body(base_url, api_client):
+    # Act
+    response = api_client.post(f"{base_url}/login", json={})
+
+    # Assert
+    assert response.status_code == 422
+
+
+def test_login_valid_credentials(base_url, api_client):
+    # Arrange
+    payload = {"email": "test@azaion.com", "password": "testpass"}
+
+    # Act
+    response = api_client.post(f"{base_url}/login", json=payload)
+
+    # Assert
+    assert response.status_code == 200
+    assert response.json()["status"] == "ok"
+
+
+def test_status_authenticated_after_login(base_url, logged_in_client):
+    # Act
+    response = logged_in_client.get(f"{base_url}/status")
+
+    # Assert
+    assert response.status_code == 200
+    assert response.json()["authenticated"] is True

e2e/tests/test_health.py

@@ -0,0 +1,7 @@
+def test_health_returns_200(base_url, api_client):
+    # Act
+    response = api_client.get(f"{base_url}/health")
+
+    # Assert
+    assert response.status_code == 200
+    assert response.json()["status"] == "healthy"

e2e/tests/test_performance.py

@@ -0,0 +1,17 @@
+import time
+
+
+def test_health_latency_p95(base_url, api_client):
+    # Arrange
+    times = []
+    # Act
+    for _ in range(100):
+        start = time.perf_counter()
+        response = api_client.get(f"{base_url}/health")
+        elapsed = time.perf_counter() - start
+        times.append(elapsed)
+        response.raise_for_status()
+    times.sort()
+    p95 = times[94]
+    # Assert
+    assert p95 <= 0.1

e2e/tests/test_resources.py

@@ -0,0 +1,74 @@
+import pytest
+
+
+def test_upload_resource(base_url, logged_in_client):
+    # Arrange
+    url = f"{base_url}/upload/testmodel"
+    files = {"data": ("testmodel.bin", b"test content")}
+    data = {"folder": "models"}
+
+    # Act
+    response = logged_in_client.post(url, files=files, data=data)
+
+    # Assert
+    assert response.status_code == 200
+    assert response.json()["status"] == "ok"
+
+
+def test_download_resource(base_url, logged_in_client):
+    # Arrange
+    url = f"{base_url}/load/testmodel"
+    body = {"filename": "testmodel", "folder": "models"}
+
+    # Act
+    response = logged_in_client.post(url, json=body)
+
+    # Assert
+    assert response.status_code == 200
+    assert len(response.content) > 0
+
+
+def test_download_nonexistent(base_url, logged_in_client):
+    # Arrange
+    url = f"{base_url}/load/nonexistent"
+    body = {"filename": "nonexistent", "folder": "nonexistent"}
+
+    # Act
+    response = logged_in_client.post(url, json=body)
+
+    # Assert
+    assert response.status_code == 500
+
+
+def test_upload_no_file(base_url, logged_in_client):
+    # Arrange
+    url = f"{base_url}/upload/testfile"
+
+    # Act
+    response = logged_in_client.post(url, data={"folder": "models"})
+
+    # Assert
+    assert response.status_code == 422
+
+
+def test_upload_download_roundtrip(base_url, logged_in_client):
+    # Arrange
+    filename = "roundtrip"
+    folder = "models"
+    content = b"roundtrip-payload-data"
+    upload_url = f"{base_url}/upload/{filename}"
+    load_url = f"{base_url}/load/{filename}"
+    files = {"data": (f"{filename}.bin", content)}
+    data = {"folder": folder}
+
+    # Act
+    upload_response = logged_in_client.post(upload_url, files=files, data=data)
+    download_response = logged_in_client.post(
+        load_url,
+        json={"filename": filename, "folder": folder},
+    )
+
+    # Assert
+    assert upload_response.status_code == 200
+    assert download_response.status_code == 200
+    assert download_response.content == content

e2e/tests/test_unlock.py

@@ -0,0 +1,66 @@
+import os
+import subprocess
+import time
+
+
+COMPOSE_FILE = os.path.join(os.path.dirname(__file__), "..", "docker-compose.test.yml")
+
+
+def _compose_exec(cmd: str):
+    subprocess.run(
+        ["docker", "compose", "-f", COMPOSE_FILE, "exec", "system-under-test",
+         "bash", "-c", cmd],
+        capture_output=True, timeout=15,
+    )
+
+
+def _wait_for_settled(base_url, client, timeout=30):
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        resp = client.get(f"{base_url}/unlock/status")
+        state = resp.json()["state"]
+        if state in ("idle", "error", "ready"):
+            return state
+        time.sleep(0.5)
+    return None
+
+
+def test_unlock_status_idle(base_url, api_client):
+    # Act
+    response = api_client.get(f"{base_url}/unlock/status")
+
+    # Assert
+    assert response.status_code == 200
+    data = response.json()
+    assert data["state"] == "idle"
+    assert data["error"] is None
+
+
+def test_unlock_missing_archive(base_url, api_client):
+    # Arrange
+    payload = {"email": "test@azaion.com", "password": "testpass"}
+
+    # Act
+    response = api_client.post(f"{base_url}/unlock", json=payload)
+
+    # Assert
+    assert response.status_code == 404
+
+
+def test_unlock_concurrent_returns_current_state(base_url, api_client):
+    # Arrange
+    _compose_exec("dd if=/dev/urandom of=/tmp/test.enc bs=1024 count=1 2>/dev/null")
+    payload = {"email": "test@azaion.com", "password": "testpass"}
+
+    try:
+        # Act
+        first = api_client.post(f"{base_url}/unlock", json=payload)
+        second = api_client.post(f"{base_url}/unlock", json=payload)
+
+        # Assert
+        assert first.status_code == 200
+        assert second.status_code == 200
+        assert second.json()["state"] != "idle"
+    finally:
+        _compose_exec("rm -f /tmp/test.enc /tmp/test.tar")
+        _wait_for_settled(base_url, api_client)

e2e/tests/test_zz_resilience.py

@@ -0,0 +1,72 @@
+import os
+import subprocess
+import time
+
+
+COMPOSE_FILE = os.path.join(os.path.dirname(__file__), "..", "docker-compose.test.yml")
+
+
+def _compose(*args):
+    subprocess.run(
+        ["docker", "compose", "-f", COMPOSE_FILE, *args],
+        capture_output=True, timeout=30,
+    )
+
+
+def test_download_when_cdn_unavailable(base_url, logged_in_client):
+    # Arrange
+    _compose("stop", "mock-cdn")
+    time.sleep(1)
+
+    try:
+        # Act
+        try:
+            response = logged_in_client.post(
+                f"{base_url}/load/nocache",
+                json={"filename": "nocache", "folder": "models"},
+                timeout=15,
+            )
+            status = response.status_code
+        except Exception:
+            status = 0
+
+        # Assert
+        assert status != 200
+    finally:
+        _compose("start", "mock-cdn")
+        time.sleep(3)
+
+
+def test_unlock_with_corrupt_archive(base_url, api_client):
+    # Arrange
+    subprocess.run(
+        ["docker", "compose", "-f", COMPOSE_FILE, "exec", "system-under-test",
+         "bash", "-c", "dd if=/dev/urandom of=/tmp/test.enc bs=1024 count=1 2>/dev/null"],
+        capture_output=True, timeout=15,
+    )
+    payload = {"email": "test@azaion.com", "password": "testpass"}
+
+    try:
+        # Act
+        response = api_client.post(f"{base_url}/unlock", json=payload)
+        assert response.status_code == 200
+
+        deadline = time.monotonic() + 30
+        body = None
+        while time.monotonic() < deadline:
+            status = api_client.get(f"{base_url}/unlock/status")
+            body = status.json()
+            if body["state"] in ("error", "ready"):
+                break
+            time.sleep(0.5)
+
+        # Assert
+        assert body is not None
+        assert body["state"] == "error"
+        assert body["error"] is not None
+    finally:
+        subprocess.run(
+            ["docker", "compose", "-f", COMPOSE_FILE, "exec", "system-under-test",
+             "bash", "-c", "rm -f /tmp/test.enc /tmp/test.tar"],
+            capture_output=True, timeout=15,
+        )