def METHOD_NAME(
    long_log_with_features, user_features, item_features, two_stages_kwargs,
):
    two_stages = TwoStagesScenario(**two_stages_kwargs)
    two_stages.fit(
        long_log_with_features,
        user_features,
        item_features.filter(sf.col("iq") > 4),
    )
    pred = two_stages.predict(
        log=long_log_with_features,
        k=2,
        user_features=user_features,
        item_features=item_features,
    )
    assert pred.count() == 6
    assert sorted(pred.select(sf.collect_set("user_idx")).collect()[0][0]) == [
        0,
        1,
        2,
    ]