import types

from embrace.mapobject import mapobject
from embrace.mapobject import RowMapper
from embrace.query import joinone


class TestRowMapper:
    def test_get_split_points_by_split(self):
        rm = RowMapper([mapobject(tuple), mapobject(tuple, split="c")])
        rm.set_description(["a", "b", "c", "d"])  # type: ignore
        assert rm.get_split_points() == [slice(0, 2), slice(2, 4)]

    def test_get_split_points_by_count(self):
        rm = RowMapper([mapobject(tuple, column_count=3), mapobject(tuple)])
        rm.set_description(["a", "b", "c", "d"])  # type: ignore
        assert rm.get_split_points() == [slice(0, 3), slice(3, 4)]

    def test_row_spec_after_joins_returns_row_spec(self):
        A = types.new_class("A")
        B = types.new_class("B")
        ma = mapobject(A)
        mb = mapobject(B)
        assert RowMapper([ma, mb]).row_spec_after_joins() == [ma, mb]

    def test_row_spec_after_joins_removes_join_source(self):
        A = types.new_class("A")
        B = types.new_class("B")
        ma = mapobject(A)
        mb = mapobject(B)
        assert RowMapper(
            [ma, mb], joins=[joinone(A, "a", B)]
        ).row_spec_after_joins() == [ma]

    def test_row_spec_after_joins_removes_labelled_join_source(self):
        A = types.new_class("A")
        B = types.new_class("B")
        ma = mapobject(A, label="classA")
        mb = mapobject(B, label="classB")
        assert RowMapper(
            [ma, mb], joins=[joinone("classA", "a", "classB")]
        ).row_spec_after_joins() == [ma]

    def test_row_spec_after_joins_removes_nested_join_source(self):
        A = types.new_class("A")
        B = types.new_class("B")
        C = types.new_class("C")
        ma = mapobject(A)
        mb = mapobject(B)
        mc = mapobject(B)
        assert RowMapper(
            [ma, mb, mc], joins=[joinone(A, "a", B), joinone(B, "b", C)]
        ).row_spec_after_joins() == [ma]
