import unittest

from autark import JsonReader, OnceFrame, P


class TestAsofJoin(unittest.TestCase):
    def test_asof_backward_join_shape(self) -> None:
        left_rows = [
            {"ts": 1, "left_v": 10.0},
            {"ts": 3, "left_v": 20.0},
            {"ts": 6, "left_v": 30.0},
        ]
        right_rows = [
            {"ts": 2, "right_v": 100.0},
            {"ts": 5, "right_v": 500.0},
        ]

        frame = OnceFrame([JsonReader(left_rows), JsonReader(right_rows)])
        p = P()
        left = p.columns(None, index=0)
        right = p.columns(None, index=1)

        left.join(
            right,
            left_on=left.columns("ts"),
            right_on=right.columns("ts"),
            kind="asof_backward",
        ).alias("asof_backward")

        realized = frame.realize(p)
        table = realized["asof_backward"]

        self.assertTrue(table)
        first_column = next(iter(table.values()))
        self.assertEqual(len(first_column), len(left_rows))


if __name__ == "__main__":
    unittest.main()
