fix(gain_loss): list masks crash, a non-unique index gives a false overlap error, and the length check is unreachable #595

import pandas as pd
import pytest

from openretailscience.analysis.gain_loss import GainLoss


@pytest.fixture
def cola_df() -> pd.DataFrame:
    """Coke and Pepsi transactions across 2023 and 2024 for five customers."""
    return pd.DataFrame(
        {
            "customer_id": [101, 101, 102, 102, 103, 103, 104, 105],
            "brand": ["Coke", "Pepsi", "Coke", "Coke", "Pepsi", "Coke", "Pepsi", "Coke"],
            "year": [2023, 2024, 2023, 2024, 2023, 2024, 2024, 2023],
            "unit_spend": [10.0, 12.0, 8.0, 5.0, 7.0, 9.0, 4.0, 6.0],
        },
    )


def _masks(df: pd.DataFrame) -> dict[str, pd.Series]:
    return {
        "p1_index": df["year"] == 2023,
        "p2_index": df["year"] == 2024,
        "focus_group_index": df["brand"] == "Coke",
        "comparison_group_index": df["brand"] == "Pepsi",
    }


def _gain_loss(df: pd.DataFrame, masks: dict, **kwargs: str) -> GainLoss:
    return GainLoss(df=df, focus_group_name="Coke", comparison_group_name="Pepsi", **masks, **kwargs)


class TestMaskInputs:
    def test_list_masks_match_series_masks(self, cola_df):
        """Python lists of bools give the same decomposition as boolean Series."""
        masks = _masks(cola_df)
        expected = _gain_loss(cola_df, masks).gain_loss_df
        result = _gain_loss(cola_df, {name: mask.tolist() for name, mask in masks.items()}).gain_loss_df
        pd.testing.assert_frame_equal(result, expected)

    def test_non_unique_index_with_disjoint_masks_is_accepted(self, cola_df):
        """Disjoint masks on a frame indexed by repeated customer ids are not reported as overlapping."""
        expected = _gain_loss(cola_df, _masks(cola_df)).gain_loss_df
        indexed = cola_df.copy()
        indexed.index = indexed["customer_id"].to_numpy()
        result = _gain_loss(indexed, _masks(indexed)).gain_loss_df
        pd.testing.assert_frame_equal(result, expected)

    def test_overlapping_masks_still_raise_with_non_unique_index(self, cola_df):
        """Masks that truly overlap are still rejected."""
        indexed = cola_df.copy()
        indexed.index = indexed["customer_id"].to_numpy()
        masks = _masks(indexed)
        masks["p2_index"] = indexed["year"] >= 2023
        with pytest.raises(ValueError, match="p1_index and p2_index should not overlap"):
            _gain_loss(indexed, masks)


class TestMaskLengthValidation:
    @pytest.mark.parametrize("as_list", [True, False], ids=["list", "series"])
    def test_short_mask_raises_same_length_error(self, cola_df, as_list):
        """The documented same-length ValueError is raised for list and Series masks."""
        masks = _masks(cola_df)
        masks["p1_index"] = masks["p1_index"].iloc[:-1]
        if as_list:
            masks["p1_index"] = masks["p1_index"].tolist()
        with pytest.raises(ValueError, match="should have the same length"):
            _gain_loss(cola_df, masks)