import remedapy as R


class TestFlatMap:
    def test_data_first(self):
        # R.flat_map(data, callbackfn)
        assert list(R.flat_map([1, 2, 3], lambda x: [x, x * 10])) == [1, 10, 2, 20, 3, 30]
        assert list(R.flat_map([1, 2, 3], lambda: [0, 1])) == [0, 1, 0, 1, 0, 1]
        assert list(R.flat_map([1, 2, 3], lambda x, i: [x, x * 10 + i])) == [1, 10, 2, 21, 3, 32]
        assert list(R.flat_map([1, 2, 3], lambda x, i, arr: [x + i, x * 10 + len(arr)])) == [1, 13, 3, 23, 5, 33]

    def test_data_last(self):
        # R.flat_map(callbackfn)(data)
        def fn(x: int):
            return [x, x * 10]

        assert R.pipe([1, 2, 3], R.flat_map(fn), list) == [1, 10, 2, 20, 3, 30]
