unstructured/test_unstructured_ingest/unit/test_enhanced_dataclass.py

import json
from dataclasses import Field, dataclass, fields

import pytest

from unstructured.ingest.enhanced_dataclass import EnhancedDataClassJsonMixin, enhanced_field
from unstructured.ingest.enhanced_dataclass.dataclasses import EnhancedField


@dataclass
class AuthData(EnhancedDataClassJsonMixin):
    username: str
    password: str = enhanced_field(sensitive=True)
    date: int = enhanced_field(overload_name="time")


auth = AuthData(username="my name", password="top secret", date=3)


def test_enhanced_field():
    fs = fields(AuthData)
    for f in fs:
        if f.name == "username":
            assert isinstance(f, Field)
            assert hasattr(f, "sensitive") is False
        else:
            assert isinstance(f, EnhancedField)
            if f.name == "password":
                assert f.sensitive is True
            else:
                assert not f.sensitive


@pytest.mark.parametrize(
    ("apply_name_overload", "expected_dict"),
    [
        (True, {"username": "my name", "password": "THIS IS REDACTED", "time": 3}),
        (False, {"username": "my name", "password": "THIS IS REDACTED", "date": 3}),
    ],
)
def test_to_json(apply_name_overload: bool, expected_dict: dict):
    j = auth.to_json(
        redact_sensitive=True,
        redacted_text="THIS IS REDACTED",
        apply_name_overload=apply_name_overload,
    )
    expected = json.dumps(expected_dict)
    assert j == expected


@pytest.mark.parametrize(
    ("apply_name_overload", "expected_dict"),
    [
        (True, {"username": "my name", "password": "***REDACTED***", "time": 3}),
        (False, {"username": "my name", "password": "***REDACTED***", "date": 3}),
    ],
)
def test_to_dict(apply_name_overload: bool, expected_dict: dict):
    d = auth.to_dict(redact_sensitive=True, apply_name_overload=apply_name_overload)
    assert d == expected_dict