tests/test_security.py

import pytest

from app.security import (
    generate_session_id, generate_worker_id, generate_job_id,
    hash_payload, hash_json, constant_time_compare,
    create_worker_join_token, verify_worker_join_token,
    validate_payload_size, validate_allowed_job_type,
    validate_origin, sanitize_public_error,
)
from app.config import Settings


def test_payload_hash():
    h1 = hash_payload(b"hello")
    h2 = hash_payload(b"hello")
    h3 = hash_payload(b"world")
    assert h1 == h2
    assert h1 != h3
    assert len(h1) == 64


def test_payload_size_limit():
    from app.config import Settings
    small = b"x" * 100
    large = b"x" * (Settings.get_max_payload_bytes() + 1)
    assert validate_payload_size(small) is True
    assert validate_payload_size(large) is False


def test_constant_time_compare():
    assert constant_time_compare("abc", "abc") is True
    assert constant_time_compare("abc", "abC") is False


def test_origin_validation():
    assert validate_origin("https://huggingface.co") is True
    assert validate_origin("http://localhost:7860") is True
    assert validate_origin(None) is True


def test_public_error_sanitization():
    assert sanitize_public_error("secret key leaked") == "Internal error"
    assert sanitize_public_error("normal error") == "normal error"


def test_valid_job_type():
    assert validate_allowed_job_type("text_embedding") is True
    assert validate_allowed_job_type("invalid_type") is False