import unittest

from autark import JsonReader, OnceFrame, P


class TestOnceFrameJson(unittest.TestCase):
    def test_slice_and_reduce_count(self) -> None:
        rows = [
            {"city": "A", "v": 1.0},
            {"city": "B", "v": 2.0},
            {"city": "C", "v": 3.0},
        ]

        frame = OnceFrame(JsonReader(rows))
        p = P()

        p.columns(None).slice(0, 2).alias("head", frame.schema())
        p.columns("v").reduce("count").alias("row_count")

        realized = frame.realize(p)

        self.assertIn("head", realized)
        self.assertIn("row_count", realized)
        row_count_table = realized["row_count"]
        self.assertTrue(row_count_table)
        row_count_values = next(iter(row_count_table.values()))
        self.assertEqual(int(row_count_values[0]), len(rows))
        self.assertEqual(len(realized["head"]["city"]), 2)


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