import unittest
from nose.tools import eq_
from helix.calculators.subarray_helper import get_subarray_sizes_and_rows, extract_subarray
from helix.models.panel import Panel
from helix.models.subarray import Subarray


class SubarrayHelperTest(unittest.TestCase):
    def test_get_subarray_sizes_and_rows(self):
        panels = [
            Panel(ballast=5, subarray=2),
            Panel(ballast=6, subarray=2),
            Panel(ballast=2, subarray=2),
            Panel(ballast=1, subarray=5),
            Panel(ballast=5, subarray=5),
            Panel(ballast=6, subarray=5),
            Panel(ballast=8, subarray=5),
            Panel(ballast=2, subarray=5),
            Panel(ballast=2, subarray=6),
            Panel(ballast=2, subarray=6)
        ]

        result = get_subarray_sizes_and_rows(panels)
        expected = [
            Subarray(subarray_number=2, start_row=0, size=3),
            Subarray(subarray_number=5, start_row=3, size=5),
            Subarray(subarray_number=6, start_row=8, size=2)
        ]

        eq_(expected, result)

    def test_extract_subarray(self):
        panels = [
            Panel(ballast=5, subarray=2),
            Panel(ballast=6, subarray=2),
            Panel(ballast=2, subarray=2),
            Panel(ballast=1, subarray=5),
            Panel(ballast=5, subarray=5),
            Panel(ballast=6, subarray=5),
            Panel(ballast=8, subarray=5),
            Panel(ballast=2, subarray=5),
            Panel(ballast=2, subarray=6),
            Panel(ballast=2, subarray=6)
        ]

        result = extract_subarray(panels, 5)

        expected = [
            Panel(ballast=1, subarray=5),
            Panel(ballast=5, subarray=5),
            Panel(ballast=6, subarray=5),
            Panel(ballast=8, subarray=5),
            Panel(ballast=2, subarray=5)
        ]
        eq_(result, expected)