import pytest
from common.src.utils.xlsx import auto_fit_columns


@pytest.fixture
def mock_worksheet(mocker):
    return mocker.Mock()


@pytest.mark.parametrize("as_iterator", [True, False])
def test_auto_fit_columns_sets_expected_widths(mock_worksheet, as_iterator):
    data = [
        ["Header1", "Header2"],
        ["Short", "MuchLongerValue"],
        ["Tiny", "123"],
        ["OhNo", None],
    ]

    if as_iterator:
        data = iter(data)

    auto_fit_columns(mock_worksheet, data)

    # Expected max lengths:
    #   col0 = len("Header1") = 7
    #   col1 = len("MuchLongerValue") = 15
    # So widths should be set to 9 and 17 (add 2 padding)
    mock_worksheet.set_column.assert_any_call(0, 0, 9)
    mock_worksheet.set_column.assert_any_call(1, 1, 17)
    assert (
        mock_worksheet.set_column.call_count == 2
    ), "set_column should be called twice, once for each column which is not empty"


@pytest.mark.parametrize("as_iterator", [True, False])
def test_auto_fit_columns_with_empty_data(mock_worksheet, as_iterator):
    data = []
    if as_iterator:
        data = iter(data)
    auto_fit_columns(mock_worksheet, data)
    assert (
        not mock_worksheet.set_column.called
    ), "set_column should not be called with empty data"