from __future__ import annotations

import unittest

from build_dataset import split_way_nodes


class BuildDatasetTest(unittest.TestCase):
    def test_junction_split_preserves_graph_connectivity(self) -> None:
        nodes = [
            (10, 48.0, 2.0),
            (11, 48.0, 2.1),
            (12, 48.0, 2.2),
            (13, 48.0, 2.3),
        ]

        parts = split_way_nodes(nodes, junctions={12})

        self.assertEqual([[10, 11, 12], [12, 13]], [[node[0] for node in part] for part in parts])
        self.assertEqual(parts[0][-1][0], parts[1][0][0])


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