DataOps-platform/tests/runner/test_artifacts.py

from __future__ import annotations

import io
from datetime import UTC, datetime
from types import SimpleNamespace

import polars as pl
import pytest

from app.core.common.identifiers import new_governance_uid


class Response(io.BytesIO):
    def release_conn(self):
        return None


class FakeMinio:
    def __init__(self):
        self.buckets = {"dataops-rules"}
        self.objects = {}

    def bucket_exists(self, bucket):
        return bucket in self.buckets

    def make_bucket(self, bucket):
        self.buckets.add(bucket)

    def put_object(
        self,
        bucket,
        key,
        data,
        length,
        *,
        content_type,
        metadata,
    ):
        payload = data.read(length)
        self.objects[(bucket, key)] = {
            "payload": payload,
            "content_type": content_type,
            "metadata": {
                f"x-amz-meta-{name.lower()}": str(value)
                for name, value in metadata.items()
            },
        }

    def stat_object(self, bucket, key):
        item = self.objects[(bucket, key)]
        return SimpleNamespace(
            size=len(item["payload"]),
            content_type=item["content_type"],
            metadata=item["metadata"],
        )

    def get_object(self, bucket, key):
        return Response(self.objects[(bucket, key)]["payload"])


def _store(client, *, clock=None, max_rows=100):
    from app.runner.artifacts import ArtifactStore

    return ArtifactStore(
        client,
        bucket="dataops-rules",
        max_artifact_bytes=1024 * 1024,
        max_rows=max_rows,
        memory_limit_bytes=4 * 1024 * 1024,
        max_ttl_seconds=3600,
        clock=clock,
    )


def test_artifact_store_generates_key_and_round_trips_digest_bound_lazyframe():
    client = FakeMinio()
    store = _store(client)
    correlation_id = new_governance_uid()

    artifact = store.write(
        pl.DataFrame(
            {
                "customer_id": [1, 2],
                "name": ["Alice", "Bob"],
            }
        ).lazy(),
        correlation_id,
        300,
    )

    assert artifact["artifact_ref"].startswith(
        f"minio://dataops-rules/rules/{correlation_id}/"
    )
    assert artifact["artifact_ref"].endswith(".parquet")
    assert artifact["digest"]
    assert artifact["row_count"] == 2
    assert artifact["schema_hash"]
    assert artifact["expires_at"].endswith("Z")
    assert "dataops-test" not in repr(artifact)
    assert store.describe(artifact["artifact_ref"]) == artifact
    frame = store.read(artifact["artifact_ref"], artifact["digest"])
    assert isinstance(frame, pl.LazyFrame)
    assert frame.collect().to_dicts() == [
        {"customer_id": 1, "name": "Alice"},
        {"customer_id": 2, "name": "Bob"},
    ]


@pytest.mark.parametrize(
    ("mutation", "message"),
    [
        (
            lambda item: item["metadata"].update(
                {"x-amz-meta-sha256": "0" * 64}
            ),
            "digest",
        ),
        (
            lambda item: item["metadata"].update(
                {"x-amz-meta-schema-sha256": "0" * 64}
            ),
            "schema",
        ),
        (
            lambda item: item["metadata"].update(
                {"x-amz-meta-expires-at": "2000-01-01T00:00:00Z"}
            ),
            "expired",
        ),
        (
            lambda item: item.update({"content_type": "text/plain"}),
            "content type",
        ),
    ],
)
def test_artifact_store_rejects_tampered_digest_schema_ttl_and_content(
    mutation, message
):
    client = FakeMinio()
    now = datetime(2026, 7, 23, 10, 0, tzinfo=UTC)
    store = _store(client, clock=lambda: now)
    artifact = store.write(
        pl.DataFrame({"id": [1]}).lazy(),
        new_governance_uid(),
        300,
    )
    key = artifact["artifact_ref"].split("/", 3)[-1]
    mutation(client.objects[("dataops-rules", key)])

    with pytest.raises(ValueError, match=message):
        store.read(artifact["artifact_ref"], artifact["digest"])


def test_artifact_store_rejects_rows_size_ttl_and_unowned_references():
    client = FakeMinio()
    store = _store(client, max_rows=2)

    with pytest.raises(ValueError, match="row"):
        store.write(
            pl.DataFrame({"id": [1, 2, 3]}).lazy(),
            new_governance_uid(),
            60,
        )
    with pytest.raises(ValueError, match="TTL"):
        store.write(
            pl.DataFrame({"id": [1]}).lazy(),
            new_governance_uid(),
            7200,
        )
    with pytest.raises(ValueError, match="artifact reference"):
        store.read(
            "minio://other-bucket/rules/unsafe/value.parquet",
            "0" * 64,
        )


def test_artifact_store_does_not_return_ref_for_corrupted_server_content():
    class CorruptingMinio(FakeMinio):
        def put_object(self, bucket, key, *args, **kwargs):
            super().put_object(bucket, key, *args, **kwargs)
            payload = self.objects[(bucket, key)]["payload"]
            self.objects[(bucket, key)]["payload"] = bytes(
                [payload[0] ^ 1]
            ) + payload[1:]

    store = _store(CorruptingMinio())

    with pytest.raises(ValueError, match="digest|size"):
        store.write(
            pl.DataFrame({"id": [1]}).lazy(),
            new_governance_uid(),
            60,
        )