Repository: ray-project/ray
File: python/ray/data/tests/datasource/test_kafka.py (Python)

Code before:
    kafka_producer.flush()
    time.sleep(1)

    with pytest.raises(ValueError, match=expected_error):
        ds = ray.data.read_kafka(
            topics=[topic],
            bootstrap_servers=[bootstrap_server],
            start_offset=start_offset,
            end_offset=end_offset,
        )
        ds.take_all()


def test_read_kafka_with_datetime_offsets(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    """Test reading Kafka messages using datetime-based start and end offsets."""
    topic = "test-datetime-offsets"

    msg_ts = datetime_to_ms(datetime(2025, 1, 15))
    time_before = datetime(2025, 1, 1)
    time_after = datetime(2025, 2, 1)

    for i in range(100):
        message = {"id": i, "value": f"message-{i}"}
        kafka_producer.send(topic, value=message, timestamp_ms=msg_ts)
    kafka_producer.flush()
    time.sleep(1)

    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset=time_before,
        end_offset=time_after,
    )
    records = ds.take_all()
    assert len(records) == 100


def test_read_kafka_datetime_partial_range(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    """Test that only messages within the datetime range are returned."""
    topic = "test-datetime-partial-range"

    # first batch at Jan 1 2025, second batch at Feb 1 2025
    batch1_ts = datetime_to_ms(datetime(2025, 1, 1))
    batch2_ts = datetime_to_ms(datetime(2025, 2, 1))
    boundary_time = datetime(2025, 1, 15)  # Between the two batches

    for i in range(50):

Review comment (bug):
If we were to only send 1-3 messages, would this test catch meaningfully fewer bugs/issues? The behavior under test is "datetime offsets correctly bound which messages are read." Feel like 100 is kind of a magic number and won't necessarily verify anything that 1-3 wouldn't

Code after:
            r"start_offset \(150\) > end_offset \(latest \(resolved to 100\)\) for partition 0 in topic test-invalid-offsets-3",
            "test-invalid-offsets-3",
        ),
    ],
)
def test_read_kafka_invalid_offsets(
    bootstrap_server,
    kafka_producer,
    ray_start_regular_shared,
    start_offset,
    end_offset,
    expected_error,
    topic,
):
    for i in range(100):
        message = {"id": i, "value": f"message-{i}"}
        kafka_producer.send(topic, value=message)
    kafka_producer.flush()
    time.sleep(0.3)

    with pytest.raises(ValueError, match=expected_error):
        ds = ray.data.read_kafka(
            topics=[topic],
            bootstrap_servers=[bootstrap_server],
            start_offset=start_offset,
            end_offset=end_offset,
        )
        ds.take_all()


def test_read_kafka_with_datetime_offsets(
    bootstrap_server, kafka_producer, ray_start_regular_shared
):
    """Test reading Kafka messages using datetime-based start and end offsets."""
    topic = "test-datetime-offsets"

    msg_ts = _datetime_to_ms(datetime(2025, 1, 15))
    time_before = datetime(2025, 1, 1)
    time_after = datetime(2025, 2, 1)

    for i in range(3):
        kafka_producer.send(topic, value={"id": i}, timestamp_ms=msg_ts)
    kafka_producer.flush()
    # Brief wait for consumer-side metadata propagation after flush()
    time.sleep(0.3)

    ds = ray.data.read_kafka(
        topics=[topic],
        bootstrap_servers=[bootstrap_server],
        start_offset=time_before,
        end_offset=time_after,