Skip to content

datarobot_genai.drmcp.test_utils.stubs.prediction_result_stub

prediction_result_stub

Stub for datarobot_predict.deployment.predict used in integration tests with DR client stubs.

StubPredictionResult

Bases: NamedTuple

Stub for datarobot_predict PredictionResult (dataframe + response_headers).

Source code in datarobot_genai/drmcp/test_utils/stubs/prediction_result_stub.py 
23
24
25
26
27
class StubPredictionResult(NamedTuple):
    """Stub for datarobot_predict PredictionResult (dataframe + response_headers)."""

    dataframe: pd.DataFrame
    response_headers: dict[str, Any]

test_create_prediction_result 

test_create_prediction_result(deployment: Any, data_frame: DataFrame, max_explanations: int | str = 0, max_ngram_explanations: int | str | None = None, threshold_high: float | None = None, threshold_low: float | None = None, time_series_type: Any = None, forecast_point: Any = None, predictions_start_date: Any = None, predictions_end_date: Any = None, passthrough_columns: str | set[str] | None = None, explanation_algorithm: str | None = None, prediction_endpoint: str | None = None, relax_known_in_advance_features_check: bool | None = None, timeout: int = 600, **kwargs: Any) -> StubPredictionResult

Stub for datarobot_predict.deployment.predict.

Returns a result with row count and columns matching integration test expectations (sales (actual)PREDICTION, FORECAST*, explanation columns; 7 rows for forecast, 14 for historical range).

Source code in datarobot_genai/drmcp/test_utils/stubs/prediction_result_stub.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def test_create_prediction_result(
    deployment: Any,
    data_frame: pd.DataFrame,
    max_explanations: int | str = 0,
    max_ngram_explanations: int | str | None = None,
    threshold_high: float | None = None,
    threshold_low: float | None = None,
    time_series_type: Any = None,
    forecast_point: Any = None,
    predictions_start_date: Any = None,
    predictions_end_date: Any = None,
    passthrough_columns: str | set[str] | None = None,
    explanation_algorithm: str | None = None,
    prediction_endpoint: str | None = None,
    relax_known_in_advance_features_check: bool | None = None,
    timeout: int = 600,
    **kwargs: Any,
) -> StubPredictionResult:
    """Stub for datarobot_predict.deployment.predict.

    Returns a result with row count and columns matching integration test expectations
    (sales (actual)_PREDICTION, FORECAST_*, explanation columns; 7 rows for forecast,
    14 for historical range).
    """
    n_in = len(data_frame)
    # Time series: integration tests expect 7 rows for forecast_point, 14 for historical range
    if forecast_point is not None:
        n_rows = 7
    elif predictions_start_date is not None and predictions_end_date is not None:
        n_rows = 14
    else:
        n_rows = n_in

    # Build output with target row count
    if n_in == 0:
        out_df = pd.DataFrame(index=range(n_rows))
    elif n_rows <= n_in:
        out_df = data_frame.iloc[:n_rows].copy()
    else:
        out_df = pd.concat([data_frame] * (n_rows // n_in + 1), ignore_index=True).iloc[:n_rows]

    # Prediction column: use "sales (actual)_PREDICTION" when input has "sales" (integration tests)
    if "sales" in data_frame.columns:
        prediction_col = "sales (actual)_PREDICTION"
    else:
        prediction_col = "stub_PREDICTION"
    if prediction_col not in out_df.columns:
        out_df[prediction_col] = [0.0] * n_rows

    # Time series: add FORECAST_POINT and FORECAST_DISTANCE when forecast params present
    if forecast_point is not None or (
        predictions_start_date is not None and predictions_end_date is not None
    ):
        if "FORECAST_POINT" not in out_df.columns:
            out_df["FORECAST_POINT"] = (
                [forecast_point] * n_rows if forecast_point is not None else [pd.NaT] * n_rows
            )
        if "FORECAST_DISTANCE" not in out_df.columns:
            out_df["FORECAST_DISTANCE"] = list(range(1, n_rows + 1))

    # Explanations: add at least one column so tests asserting "EXPLANATION" or "SHAP" in cols pass
    if max_explanations not in (0, "0"):
        if not any("EXPLANATION" in c or "SHAP" in c for c in out_df.columns):
            out_df["EXPLANATION_1"] = [0.0] * n_rows

    return StubPredictionResult(dataframe=out_df, response_headers={})