Repository: ray-project/ray
File: python/ray/data/tests/unit/test_deduping_schema.py (Python)

Code before:
        old_schema, incoming_bundle, enforce_schemas=enforce_schemas, warn=False
    )

    assert diverged

    if enforce_schemas:
        assert out_bundle.schema == pa.schema(list(old_schema) + list(incoming_schema))
    else:
        assert out_bundle.schema == old_schema


@pytest.mark.parametrize("enforce_schemas", [False, True])
def test_dedupe_schema_mismatch(enforce_schemas: bool, caplog, propagate_logs):
    old_schema = pa.schema(
        [
            pa.field("foo", pa.int32()),
            pa.field("bar", pa.string()),
            pa.field("baz", pa.bool8()),
        ]
    )
    incoming_schema = pa.schema(
        [
            pa.field("foo", pa.int64()),
            pa.field("baz", pa.bool8()),
            pa.field("qux", pa.float64()),
        ]
    )

    incoming_bundle = RefBundle([], owns_blocks=False, schema=incoming_schema)

    # Capture warnings
    with caplog.at_level(
        logging.WARNING,
        logger="ray.data._internal.execution.streaming_executor_state",
    ):
        out_bundle, diverged = dedupe_schemas_with_validation(
            old_schema, incoming_bundle, enforce_schemas=enforce_schemas, warn=True
        )

    assert diverged

    if enforce_schemas:
        assert out_bundle.schema == pa.schema(
            [
                pa.field("foo", pa.int64()),
                pa.field("bar", pa.string()),
                pa.field("baz", pa.bool8()),
                pa.field("qux", pa.float64()),
            ]
        )


Review comment (suggestion):
Can you use a multi-string `"""` for these asserts? Since changes you are making are straightfoward, it's better to test the entire msg rather than in isolation. It also makes it self-documenting as a developer I now know the "shape/format" of the warning message

Code after:
        [
            pa.field("foo", pa.int32()),
            pa.field("bar", pa.string()),
            pa.field("baz", pa.bool_()),
            pa.field("quux", pa.uint16()),
            pa.field("corge", pa.int64()),
            pa.field("grault", pa.int32()),
        ]
    )
    incoming_schema = pa.schema(
        [
            pa.field("foo", pa.int64()),
            pa.field("baz", pa.bool_()),
            pa.field("qux", pa.float64()),
            pa.field("quux", pa.uint32()),
            pa.field("corge", pa.int32()),
            pa.field("grault", pa.int64()),
        ]
    )

    msg = _build_schemas_mismatch_warning(
        old_schema, incoming_schema, truncation_length=truncation_length
    )

    if truncation_length >= 4:
        assert (
            """Operator produced a RefBundle with a different schema than the previous one.
Fields exclusive to the incoming schema (1 total):
    qux: double
Fields exclusive to the old schema (1 total):
    bar: string
Fields that have different types across the old and the incoming schemas (4 total):
    foo: int32 => int64
    quux: uint16 => uint32
    corge: int64 => int32
    grault: int32 => int64
This may lead to unexpected behavior."""
            == msg
        )
    elif truncation_length == 3:
        assert (
            """Operator produced a RefBundle with a different schema than the previous one.
Fields exclusive to the incoming schema (1 total):
    qux: double
Fields exclusive to the old schema (1 total):
    bar: string
Fields that have different types across the old and the incoming schemas (4 total):
    foo: int32 => int64
    quux: uint16 => uint32
    corge: int64 => int32
    ... and 1 more