"""Test for generator exercises"""
import unittest
from io import StringIO

from helpers import BaseTestCase
from generators import (
    join,
    strip_lines,
    all_together,
    parse_ranges,
    first_prime_over,
    get_primes_over,
    head,
)


class JoinTests(BaseTestCase):

    """Tests for join."""

    def test_strings_with_default_separator(self):
        self.assertEqual(join(["a", "b", "c"]), "a b c")
        self.assertEqual(join(["purple", "green", "blue"]), "purple green blue")

    def test_numbers_with_default_separator(self):
        self.assertEqual(join([2, 1, 3, 4, 7, 11]), "2 1 3 4 7 11")
        self.assertEqual(join([1.5, 8.4, 9.2]), "1.5 8.4 9.2")

    def test_strings_with_custom_separator(self):
        self.assertEqual(join(["a", "b", "c"], sep=", "), "a, b, c")
        self.assertEqual(
            join(["purple", "green", "blue"], sep="\n"),
            "purple\ngreen\nblue",
        )

    def test_numbers_with_custom_separator(self):
        self.assertEqual(join([2, 1, 3, 4, 7, 11], sep="."), "2.1.3.4.7.11")
        self.assertEqual(join([1.5, 8.4, 9.2], sep="---"), "1.5---8.4---9.2")

    def test_tuples(self):
        self.assertEqual(
            join([(2, 1), (3, 4), (7, 11)], sep="\n"),
            "(2, 1)\n(3, 4)\n(7, 11)",
        )
        self.assertEqual(
            join([(2, 1), (3, 4), (7, 11)], sep=", "),
            "(2, 1), (3, 4), (7, 11)",
        )


class StripLinesTests(BaseTestCase):

    """Tests for strip_lines."""

    def test_strip_lines_list(self):
        input_lines = ["line 1\n", "line 2\n"]
        expected_output = ['line 1', 'line 2']
        self.assertEqual(list(strip_lines(input_lines)), expected_output)

    def test_strip_lines_file_object(self):
        file_object = StringIO("line 1\nline 2\n")  # StringIO is a file-like object
        expected_output = ['line 1', 'line 2']
        self.assertEqual(list(strip_lines(file_object)), expected_output)

    def test_strip_lines_empty_input(self):
        input_lines = []
        expected_output = []
        self.assertEqual(list(strip_lines(input_lines)), expected_output)

    def test_strip_lines_no_newlines(self):
        input_lines = ["line 1", "line 2"]
        expected_output = ['line 1', 'line 2']
        self.assertEqual(list(strip_lines(input_lines)), expected_output)

    # To test bonus 1, comment out the next line
    @unittest.expectedFailure
    def test_returns_an_iterator(self):
        file_object = StringIO("line 1\nline 2\n")  # StringIO is a file-like object
        output = strip_lines(file_object)
        self.assertEqual(file_object.tell(), 0, "At beginning of file")
        self.assertEqual(next(output), "line 1")
        self.assertEqual(file_object.tell(), 7, "Just after line 1")
        self.assertEqual(next(output), "line 2")


class AllTogetherTests(unittest.TestCase):

    """Tests for all_together."""

    def test_list_and_tuple(self):
        outputs = list(all_together([1, 2], (3, 4)))
        expected = [1, 2, 3, 4]
        self.assertEqual(outputs, expected)

    def test_with_strings(self):
        outputs = list(all_together([1, 2], (3, 4), "hello"))
        expected = [1, 2, 3, 4, "h", "e", "l", "l", "o"]
        self.assertEqual(outputs, expected)

    def test_empty_list(self):
        outputs = list(all_together([], (), "", [1, 2]))
        self.assertEqual(outputs, [1, 2])

    def test_iterator(self):
        outputs = all_together([1], [2])
        self.assertEqual(list(outputs), [1, 2])
        self.assertEqual(list(outputs), [])


class ParseRangesTests(unittest.TestCase):

    """Tests for parse_ranges."""

    def test_three_ranges(self):
        self.assertEqual(
            list(parse_ranges("1-2,4-4,8-10")),
            [1, 2, 4, 8, 9, 10],
        )

    def test_with_spaces(self):
        self.assertEqual(
            list(parse_ranges("0-0, 4-8, 20-21, 43-45")),
            [0, 4, 5, 6, 7, 8, 20, 21, 43, 44, 45],
        )


class FirstPrimeOverTests(unittest.TestCase):

    """Tests for first_prime_over."""

    def test_first_prime_over_one_million(self):
        self.assertEqual(first_prime_over(1000000), 1000003)

    def test_first_prime_over_three_million(self):
        self.assertEqual(first_prime_over(3000000), 3000017)


class GetPrimesOverTests(unittest.TestCase):

    """Tests for get_primes_over."""

    def test_10_primes_over_one_million(self):
        ten_primes = [
            1000003,
            1000033,
            1000037,
            1000039,
            1000081,
            1000099,
            1000117,
            1000121,
            1000133,
            1000151,
        ]
        self.assertSequenceEqual(list(get_primes_over(10)), ten_primes)

    def test_first_prime_over_one_million(self):
        self.assertSequenceEqual(list(get_primes_over(1)), [1000003])


class HeadTests(unittest.TestCase):

    """Tests for head."""

    def test_first_two(self):
        self.assertEqual(list(head([1, 2, 3, 4, 5], n=2)), [1, 2])

    def test_iterator(self):
        output = head([1, 2, 3, 4, 5], n=4)
        self.assertEqual(list(zip(output, output)), [(1, 2), (3, 4)])

    def test_does_not_consume_more_than_n(self):
        numbers = [2, 1, 3, 4, 7, 11, 18, 29]
        squares = (n**2 for n in numbers)
        self.assertEqual(list(head(squares, 5)), [4, 1, 9, 16, 49])
        self.assertEqual(list(squares), [121, 324, 841])


if __name__ == "__main__":
    from helpers import error_message

    error_message()
