Skip to content

Commit

Permalink
[Tests] Add test to demonstrate new storey.Choice implementation
Browse files Browse the repository at this point in the history
  • Loading branch information
gtopper committed Sep 15, 2024
1 parent 74e8a61 commit 2f11c90
Showing 1 changed file with 54 additions and 0 deletions.
54 changes: 54 additions & 0 deletions tests/feature-store/test_steps.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import numpy as np
import pandas as pd
import pytest
import storey

import mlrun
import mlrun.feature_store as fstore
Expand Down Expand Up @@ -733,6 +734,59 @@ def test_imputer_default_value(rundb_mock, engine):
assert not imputed_df.isnull().values.any()


class MyChoice(storey.Choice):
def select_outlets(self, event):
outlets = []
r = event["routing"]
if r == "target1" or r == "both":
outlets.append("target1")
if r == "target2" or r == "both":
outlets.append("target2")
return outlets


def test_choice(rundb_mock):
df = pd.DataFrame(
{
"id": [1, 2, 3, 4],
"routing": ["target1", "target2", "both", "neither"],
}
)
feature_set = fstore.FeatureSet(
"fs-default-value",
entities=["id"],
description="feature set with nones",
)
feature_set.graph.to("MyChoice")

# Mocking
output_path = tempfile.TemporaryDirectory()
feature_set._run_db = rundb_mock
feature_set.reload = unittest.mock.Mock()
feature_set.save = unittest.mock.Mock()
feature_set.purge_targets = unittest.mock.Mock()

feature_set.ingest(
source=df,
targets=[
ParquetTarget(name="target1", path=f"{output_path.name}/target1.parquet"),
ParquetTarget(name="target2", path=f"{output_path.name}/target2.parquet"),
],
)

read_back_df1 = pd.read_parquet(feature_set.get_target_path(name="target1"))
read_back_df2 = pd.read_parquet(feature_set.get_target_path(name="target2"))

pd.testing.assert_frame_equal(
read_back_df1.reset_index(),
df.iloc[[0, 2]].reset_index(drop=True),
)
pd.testing.assert_frame_equal(
read_back_df2.reset_index(),
df.iloc[[1, 2]].reset_index(drop=True),
)


def get_data(with_none=False):
names = ["A", "B", "C", "D", "E"]
ages = [33, 4, 76, 90, 24]
Expand Down

0 comments on commit 2f11c90

Please sign in to comment.