DataOps-platform/tests/data_research/test_file_policy.py

from __future__ import annotations

import pytest


def test_extension_mime_and_magic_must_agree():
    from app.core.data_research.file_policy import FilePolicy, FilePolicyViolation

    policy = FilePolicy(max_bytes=1024)
    result = policy.validate("definition.pdf", "application/pdf", b"%PDF-1.7\nbody")
    assert result.extension == ".pdf"
    assert result.media_type == "application/pdf"

    with pytest.raises(FilePolicyViolation, match="file signature"):
        policy.validate("definition.pdf", "application/pdf", b"not-a-pdf")
    with pytest.raises(FilePolicyViolation, match="media type"):
        policy.validate("definition.csv", "image/png", b"id,name\n1,A")


def test_legacy_doc_unsafe_names_and_size_are_rejected():
    from app.core.data_research.file_policy import FilePolicy, FilePolicyViolation

    policy = FilePolicy(max_bytes=8)
    with pytest.raises(FilePolicyViolation, match="convert.*DOCX"):
        policy.validate("legacy.doc", "application/msword", b"1234")
    with pytest.raises(FilePolicyViolation, match="unsafe filename"):
        policy.validate("../secret.csv", "text/csv", b"id\n1")
    with pytest.raises(FilePolicyViolation, match="size limit"):
        policy.validate("large.csv", "text/csv", b"id\n123456")


class ArtifactRepository:
    def __init__(self):
        self.by_key = {}
        self.saved = []

    def find(self, source_uid, content_hash, parser_version):
        return self.by_key.get((source_uid, content_hash, parser_version))

    def save(self, artifact):
        key = (artifact.source_uid, artifact.content_hash, artifact.parser_version)
        self.by_key[key] = artifact
        self.saved.append(artifact)
        return artifact


class Storage:
    def __init__(self):
        self.puts = []

    def put(self, object_key, content, media_type):
        self.puts.append((object_key, content, media_type))
        return f"minio://data-research/{object_key}"


def test_artifact_storage_reuses_source_hash_and_hides_unsafe_filename():
    from app.core.data_research.artifacts import ArtifactService
    from app.core.data_research.file_policy import FilePolicy

    repository = ArtifactRepository()
    storage = Storage()
    service = ArtifactService(
        repository,
        storage,
        FilePolicy(max_bytes=1024),
        uid_factory=lambda: "artifact-1",
    )

    first, created = service.store(
        "source-1", "客户定义.csv", "text/csv", b"id,name\n1,A", "csv-v1"
    )
    second, created_again = service.store(
        "source-1", "客户定义.csv", "text/csv", b"id,name\n1,A", "csv-v1"
    )

    assert created is True and created_again is False
    assert first == second
    assert len(storage.puts) == 1
    assert "客户定义.csv" not in first.storage_ref
    assert first.content_hash in first.storage_ref