test_fn.py

#/usr/bin/env python3
# coding: utf8

from __future__ import generators
import unittest


def fn(iterator1, iterator2):
    # may be a little bit faster but not DRY if get_next() function calls will be replaced by function code
    def get_next(iterator):
        """
            @return: tuple (next value or None, is_iteration_stopped: bool)
        """
        try:
            n = next(iterator)
            return n, False
        except (StopIteration, ):
            return None, True

    n1, stop1 = get_next(iterator1)
    n2, stop2 = get_next(iterator2)

    while not (stop1 and stop2):
        if n2 is None or (n1 is not None and n1 < n2):
            yield n1
            n1, stop1 = get_next(iterator1)
        elif n2 is not None:
            yield n2
            n2, stop2 = get_next(iterator2)


class TestFn(unittest.TestCase):
    def setUp(self):
        self.arr1 = [1, 2, 3, 4, 5, 6, 7, 89]
        self.arr2 = [1, 2, 3, 4, 5, 6, 7, 23, 45, 56, 89, 90]
        self.arr3 = []

    def test_fn_arr1_to_arr2(self):
        arr = list(fn(iter(self.arr1), iter(self.arr2)))
        self.assertEqual(arr, [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 23, 45, 56, 89, 89, 90])

    def test_fn_arr2_to_arr1(self):
        arr = list(fn(iter(self.arr2), iter(self.arr1)))
        self.assertEqual(arr, [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 23, 45, 56, 89, 89, 90])

    def test_fn_arr1_to_empty(self):
        arr = list(fn(iter(self.arr1), iter(self.arr3)))
        self.assertEqual(arr, self.arr1)

    def test_fn_empty_to_arr1(self):
        arr = list(fn(iter(self.arr3), iter(self.arr1)))
        self.assertEqual(arr, self.arr1)

    def test_big_arrays(self):
        arr1 = []
        arr2 = []
        for i in range(1000000):
            arr1.append(i)
            arr2.append(i)
        arr = list(fn(iter(arr1), iter(arr2)))
        self.assertEqual(len(arr), 2000000)
        n_prev = 0
        for n in arr:
            self.assertLessEqual(n_prev, n)
            n_prev = n


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