DataOps-platform/tests/runner/test_task_tokens.py

import pytest

from app.runner.auth import (
    TaskTokenExpired,
    TaskTokenInvalid,
    TaskTokenIssuer,
    TaskTokenVerifier,
    node_digest,
)

NODE = {
    "id": "read_orders",
    "type": "sql.query",
    "data_source_uid": "01900000-0000-7000-8000-000000000010",
    "purpose": "read",
    "config": {
        "statement": "SELECT * FROM orders WHERE day = :day",
        "parameters": {"day": "${parameters.day}"},
    },
}


def test_task_token_is_short_lived_and_bound_to_one_node():
    def clock():
        return 1_000

    issuer = TaskTokenIssuer("x" * 32, clock=clock, ttl_seconds=60)
    verifier = TaskTokenVerifier("x" * 32, clock=clock)

    token = issuer.issue(
        task_uid="01900000-0000-7000-8000-000000000011",
        dataflow_uid="01900000-0000-7000-8000-000000000012",
        workflow_version=7,
        correlation_id="01900000-0000-7000-8000-000000000013",
        node=NODE,
    )
    claims = verifier.verify(token, node=NODE)

    assert claims.node_id == "read_orders"
    assert claims.node_type == "sql.query"
    assert claims.purpose == "read"
    assert claims.node_digest == node_digest(NODE)
    assert claims.expires_at == 1_060
    assert claims.jti
    assert "password" not in token.lower()


def test_task_token_rejects_expiry_tampering_and_other_node():
    now = [1_000]
    issuer = TaskTokenIssuer("x" * 32, clock=lambda: now[0], ttl_seconds=10)
    verifier = TaskTokenVerifier("x" * 32, clock=lambda: now[0])
    token = issuer.issue(
        task_uid="01900000-0000-7000-8000-000000000011",
        dataflow_uid="01900000-0000-7000-8000-000000000012",
        workflow_version=7,
        correlation_id="01900000-0000-7000-8000-000000000013",
        node=NODE,
    )

    changed = {**NODE, "config": {"statement": "DELETE FROM orders"}}
    with pytest.raises(TaskTokenInvalid, match="node binding"):
        verifier.verify(token, node=changed)

    header, payload, signature = token.split(".")
    changed_signature = ("A" if signature[0] != "A" else "B") + signature[1:]
    with pytest.raises(TaskTokenInvalid):
        verifier.verify(".".join((header, payload, changed_signature)), node=NODE)

    now[0] = 1_011
    with pytest.raises(TaskTokenExpired):
        verifier.verify(token, node=NODE)


def test_task_token_secret_has_a_minimum_strength():
    with pytest.raises(ValueError, match="at least 32"):
        TaskTokenIssuer("too-short")