from ldnode.calculator.plans.polars.window_slicer import extract_max_window


def test_extract_max_window_reads_polars_rolling_window_size_keyword():
    computed_factors = [[{
        "factor": "volatility",
        "expr": 'col("pct_chg").rolling_std(window_size=242, min_samples=1, ddof=0).over("code") * (242 ** 0.5)',
        "expr_type": 1,
        "is_intermediate": False,
    }]]

    assert extract_max_window(computed_factors) == 242


def test_extract_max_window_reads_polars_rolling_window_size_positional():
    computed_factors = [[{
        "factor": "amount_20",
        "expr": 'col("amount").rolling_sum(20, min_samples=1).over("code")',
        "expr_type": 1,
        "is_intermediate": False,
    }]]

    assert extract_max_window(computed_factors) == 20


def test_extract_max_window_keeps_stage_dependency_sum():
    computed_factors = [
        [{
            "factor": "volatility",
            "expr": 'col("pct_chg").rolling_std(window_size=242, min_samples=1, ddof=0).over("code")',
            "expr_type": 1,
            "is_intermediate": False,
        }],
        [{
            "factor": "volatility_ma",
            "expr": "ts_mean(volatility, 10)",
            "expr_type": 2,
            "is_intermediate": False,
        }],
    ]

    assert extract_max_window(computed_factors) == 252
