"""Unit tests for schemas module."""

from __future__ import annotations

import pyarrow

from arrakis.channel import Channel
from arrakis.schemas import count, metadata, publish, stream


def test_metadata_default():
    schema = metadata()
    names = schema.names
    assert "channel" in names
    assert "data_type" in names
    assert "publisher" in names
    assert "stride" in names
    assert "max_latency" in names


def test_metadata_partition_scope():
    schema = metadata(scope="partition")
    names = schema.names
    assert "channel" in names
    assert "data_type" in names
    assert "publisher" not in names
    assert "stride" not in names
    assert "max_latency" not in names


def test_count_schema():
    schema = count()
    assert schema.names == ["count"]
    assert schema.field("count").type == pyarrow.int64()


def test_stream_schema():
    channels = [
        Channel(name="H1:TEST-A", data_type="float64", sample_rate=16384),
        Channel(name="H1:TEST-B", data_type="float32", sample_rate=512),
    ]
    schema = stream(channels)
    assert schema.names[0] == "time"
    assert schema.field("time").type == pyarrow.int64()
    assert schema.field("H1:TEST-A").type == pyarrow.list_(pyarrow.float64())
    assert schema.field("H1:TEST-B").type == pyarrow.list_(pyarrow.float32())
    assert schema.field("H1:TEST-A").metadata[b"rate"] == b"16384"


def test_publish_schema():
    schema = publish()
    assert schema.names == ["properties"]
    assert schema.field("properties").type == pyarrow.map_(
        pyarrow.string(), pyarrow.string()
    )
