def secure_core_server(self, server_credentials):
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
Core.add_CoreServiceServicer_to_server(CoreServicer(), server)
server.add_secure_port("[::]:50053", server_credentials)
server.start()
yield server
server.stop(0)
def server(self):
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
Core.add_CoreServiceServicer_to_server(CoreServicer(), server)
server.add_insecure_port("[::]:50051")
server.start()
yield server
server.stop(0)
def secure_core_server(self, server_credentials):
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
Core.add_CoreServiceServicer_to_server(CoreServicer(), server)
port = find_free_port()
server.add_secure_port(f"[::]:{port}", server_credentials)
server.start()
yield port
server.stop(0)
def insecure_core_server_that_blocks_auth(self, server_credentials):
server = grpc.server(
futures.ThreadPoolExecutor(max_workers=10),
interceptors=(DisallowAuthInterceptor(), ),
)
Core.add_CoreServiceServicer_to_server(CoreServicer(), server)
server.add_insecure_port("[::]:50057")
server.start()
yield server
server.stop(0)
def serve():
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
Core.add_CoreServiceServicer_to_server(CoreServicer(), server)
server.add_insecure_port("[::]:50051")
server.start()
try:
while True:
time.sleep(_ONE_DAY_IN_SECONDS)
except KeyboardInterrupt:
server.stop(0)
def insecure_core_server_with_auth(self, server_credentials):
server = grpc.server(
futures.ThreadPoolExecutor(max_workers=10),
interceptors=(AllowAuthInterceptor(), ),
)
Core.add_CoreServiceServicer_to_server(CoreServicer(), server)
port = find_free_port()
server.add_insecure_port(f"[::]:{port}")
server.start()
yield port
server.stop(0)
def test_apply_multiple(self, client, mocker):
my_storage = Storage(id="TEST", type="redis")
my_feature_group = FeatureGroup(id="test")
my_entity = Entity(name="test")
grpc_stub = core.CoreServiceStub(grpc.insecure_channel(""))
mocker.patch.object(
grpc_stub,
'ApplyStorage',
return_value=CoreServiceTypes.ApplyStorageResponse(
storageId="TEST"))
mocker.patch.object(
grpc_stub,
'ApplyFeatureGroup',
return_value=CoreServiceTypes.ApplyFeatureGroupResponse(
featureGroupId="test"))
mocker.patch.object(
grpc_stub,
'ApplyEntity',
return_value=CoreServiceTypes.ApplyEntityResponse(
entityName="test"))
client._core_service_stub = grpc_stub
ids = client.apply([my_storage, my_entity, my_feature_group])
assert ids == ["TEST", "test", "test"]
def test_version(self, mock_client, mocker):
mock_client._core_service_stub = Core.CoreServiceStub(grpc.insecure_channel(""))
mock_client._serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel("")
)
mocker.patch.object(
mock_client._core_service_stub,
"GetFeastCoreVersion",
return_value=GetFeastCoreVersionResponse(version="0.3.2"),
)
mocker.patch.object(
mock_client._serving_service_stub,
"GetFeastServingInfo",
return_value=GetFeastServingInfoResponse(version="0.3.2"),
)
status = mock_client.version()
assert (
status["core"]["url"] == CORE_URL
and status["core"]["version"] == "0.3.2"
and status["serving"]["url"] == SERVING_URL
and status["serving"]["version"] == "0.3.2"
)
def test_list_features(self, test_client, mocker):
mocker.patch.object(
test_client,
"_core_service_stub",
return_value=Core.CoreServiceStub(grpc.insecure_channel("")),
)
feature1_proto = FeatureSpecProto(
name="feature_1", value_type=ValueProto.ValueType.FLOAT)
feature2_proto = FeatureSpecProto(
name="feature_2", value_type=ValueProto.ValueType.STRING)
mocker.patch.object(
test_client._core_service_stub,
"ListFeatures",
return_value=ListFeaturesResponse(
features={
"driver_car:feature_1": feature1_proto,
"driver_car:feature_2": feature2_proto,
}),
)
features = test_client.list_features_by_ref(project="test")
assert len(features) == 2
native_feature_list = []
for _, feature_proto in features.items():
native_feature_list.append(feature_proto)
assert sorted(native_feature_list) == sorted([
Feature.from_proto(feature1_proto),
Feature.from_proto(feature2_proto)
])
def test_ingest_dataframe_partition(self, mocked_client, mocker,
partitioned_df, tmp_path):
"""
Test ingestion with local FileSource, using DataFrame.
Partition column stated but not provided in Dataset.
"""
mocked_client._core_service_stub = Core.CoreServiceStub(
grpc.insecure_channel(""))
mocker.patch.object(
mocked_client._core_service_stub,
"GetFeatureTable",
return_value=_ingest_test_getfeaturetable_mocked_resp(
f"file://{tmp_path}", "date"),
)
mocked_client.set_project("my_project")
ft = mocked_client.get_feature_table("ingest_featuretable")
mocked_client.ingest(ft, partitioned_df, timeout=600)
pq_df = pq.read_table(tmp_path).to_pandas().drop(columns=["date"])
partitioned_df, pq_df = _ingest_test_format_dataframes(
partitioned_df, pq_df, True)
assert_frame_equal(partitioned_df, pq_df)
def test_ingest_dataframe_no_partition(self, mocked_client, mocker,
non_partitioned_df, tmp_path):
"""
Test ingestion with local FileSource, using DataFrame.
Partition column not stated.
"""
mocked_client._core_service_stub = Core.CoreServiceStub(
grpc.insecure_channel(""))
mocker.patch.object(
mocked_client._core_service_stub,
"GetFeatureTable",
return_value=_ingest_test_getfeaturetable_mocked_resp(
f"file://{tmp_path}"),
)
mocked_client.set_project("my_project")
ft = mocked_client.get_feature_table("ingest_featuretable")
mocked_client.ingest(ft, non_partitioned_df, timeout=600)
# Since not partitioning, we're only looking for single file
single_file = [
f for f in os.listdir(tmp_path)
if os.path.isfile(os.path.join(tmp_path, f))
][0]
pq_df = pq.read_table(tmp_path / single_file).to_pandas()
non_partitioned_df, pq_df = _ingest_test_format_dataframes(
non_partitioned_df, pq_df)
assert_frame_equal(non_partitioned_df, pq_df)
def test_ingest_csv(self, mocked_client, mocker, tmp_path):
"""
Test ingestion with local FileSource, using CSV file.
Partition column is provided.
"""
mocked_client._core_service_stub = Core.CoreServiceStub(
grpc.insecure_channel(""))
mocker.patch.object(
mocked_client._core_service_stub,
"GetFeatureTable",
return_value=_ingest_test_getfeaturetable_mocked_resp(
f"file://{tmp_path}", "date"),
)
partitioned_df = pd.read_csv(
os.path.join(
os.path.dirname(os.path.realpath(__file__)),
"./data/dev_featuretable.csv",
),
parse_dates=["datetime"],
)
mocked_client.set_project("my_project")
ft = mocked_client.get_feature_table("ingest_featuretable")
mocked_client.ingest(ft, partitioned_df, timeout=600)
pq_df = pq.read_table(tmp_path).to_pandas().drop(columns=["date"])
partitioned_df, pq_df = _ingest_test_format_dataframes(
partitioned_df, pq_df, True)
assert_frame_equal(partitioned_df, pq_df)
def test_get_feature_set(self, mocked_client, mocker):
mocked_client._core_service_stub = Core.CoreServiceStub(
grpc.insecure_channel("")
)
from google.protobuf.duration_pb2 import Duration
mocker.patch.object(
mocked_client._core_service_stub,
"GetFeatureSet",
return_value=GetFeatureSetResponse(
feature_set=FeatureSetProto(
spec=FeatureSetSpecProto(
name="my_feature_set",
max_age=Duration(seconds=3600),
labels={"key1": "val1", "key2": "val2"},
features=[
FeatureSpecProto(
name="my_feature_1",
value_type=ValueProto.ValueType.FLOAT,
),
FeatureSpecProto(
name="my_feature_2",
value_type=ValueProto.ValueType.FLOAT,
),
],
entities=[
EntitySpecProto(
name="my_entity_1",
value_type=ValueProto.ValueType.INT64,
)
],
source=Source(
type=SourceType.KAFKA,
kafka_source_config=KafkaSourceConfig(
bootstrap_servers="localhost:9092", topic="topic"
),
),
),
meta=FeatureSetMetaProto(),
)
),
)
mocked_client.set_project("my_project")
feature_set = mocked_client.get_feature_set("my_feature_set")
assert (
feature_set.name == "my_feature_set"
and "key1" in feature_set.labels
and feature_set.labels["key1"] == "val1"
and "key2" in feature_set.labels
and feature_set.labels["key2"] == "val2"
and feature_set.fields["my_feature_1"].name == "my_feature_1"
and feature_set.fields["my_feature_1"].dtype == ValueType.FLOAT
and feature_set.fields["my_entity_1"].name == "my_entity_1"
and feature_set.fields["my_entity_1"].dtype == ValueType.INT64
and len(feature_set.features) == 2
and len(feature_set.entities) == 1
)
def test_list_feature_sets(self, mocked_client, mocker):
mocker.patch.object(
mocked_client,
"_core_service_stub",
return_value=Core.CoreServiceStub(grpc.insecure_channel("")),
)
feature_set_1_proto = FeatureSetProto(
spec=FeatureSetSpecProto(
project="test",
name="driver_car",
max_age=Duration(seconds=3600),
labels={"key1": "val1", "key2": "val2"},
features=[
FeatureSpecProto(
name="feature_1", value_type=ValueProto.ValueType.FLOAT
)
],
)
)
feature_set_2_proto = FeatureSetProto(
spec=FeatureSetSpecProto(
project="test",
name="driver_ride",
max_age=Duration(seconds=3600),
labels={"key1": "val1"},
features=[
FeatureSpecProto(
name="feature_1", value_type=ValueProto.ValueType.FLOAT
)
],
)
)
mocker.patch.object(
mocked_client._core_service_stub,
"ListFeatureSets",
return_value=ListFeatureSetsResponse(
feature_sets=[feature_set_1_proto, feature_set_2_proto]
),
)
feature_sets = mocked_client.list_feature_sets(labels={"key1": "val1"})
assert len(feature_sets) == 2
feature_set = feature_sets[0]
assert (
feature_set.name == "driver_car"
and "key1" in feature_set.labels
and feature_set.labels["key1"] == "val1"
and "key2" in feature_set.labels
and feature_set.labels["key2"] == "val2"
and feature_set.fields["feature_1"].name == "feature_1"
and feature_set.fields["feature_1"].dtype == ValueType.FLOAT
and len(feature_set.features) == 1
)
def test_apply_single_feature_group(self, client, mocker):
my_feature_group = FeatureGroup(id="test")
grpc_stub = core.CoreServiceStub(grpc.insecure_channel(""))
with mocker.patch.object(grpc_stub, 'ApplyFeatureGroup',
return_value=CoreServiceTypes.ApplyFeatureGroupResponse(
featureGroupId="test")):
client._core_service_stub = grpc_stub
name = client.apply(my_feature_group)
assert name == "test"
def test_apply_single_entity(self, client, mocker):
my_entity = Entity(name="test")
grpc_stub = core.CoreServiceStub(grpc.insecure_channel(""))
with mocker.patch.object(grpc_stub, 'ApplyEntity',
return_value=CoreServiceTypes.ApplyEntityResponse(
entityName="test")):
client._core_service_stub = grpc_stub
name = client.apply(my_entity)
assert name == "test"
def test_apply_single_storage(self, client, mocker):
my_storage = Storage(id="TEST", type="redis")
grpc_stub = core.CoreServiceStub(grpc.insecure_channel(""))
with mocker.patch.object(grpc_stub, 'ApplyStorage',
return_value=CoreServiceTypes.ApplyStorageResponse(
storageId="TEST")):
client._core_service_stub = grpc_stub
name = client.apply(my_storage)
assert name == "TEST"
def test_apply_single_feature(self, client, mocker):
my_feature = Feature(name="test",
entity="test", granularity=Granularity.NONE)
grpc_stub = core.CoreServiceStub(grpc.insecure_channel(""))
with mocker.patch.object(grpc_stub, 'ApplyFeature',
return_value=CoreServiceTypes.ApplyFeatureResponse(
featureId="test.none.test")):
client._core_service_stub = grpc_stub
id = client.apply(my_feature)
assert id == "test.none.test"
def test_list_ingest_jobs(self, mocked_client, mocker):
mocker.patch.object(
mocked_client,
"_core_service_stub",
return_value=Core.CoreServiceStub(grpc.insecure_channel("")),
)
feature_set_proto = FeatureSetProto(
spec=FeatureSetSpecProto(
project="test", name="driver", max_age=Duration(seconds=3600),
)
)
mocker.patch.object(
mocked_client._core_service_stub,
"ListIngestionJobs",
return_value=ListIngestionJobsResponse(
jobs=[
IngestJobProto(
id="kafka-to-redis",
external_id="job-2222",
status=IngestionJobStatus.RUNNING,
feature_sets=[feature_set_proto],
source=Source(
type=SourceType.KAFKA,
kafka_source_config=KafkaSourceConfig(
bootstrap_servers="localhost:9092", topic="topic"
),
),
store=Store(name="redis"),
)
]
),
)
# list ingestion jobs by target feature set reference
ingest_jobs = mocked_client.list_ingest_jobs(
feature_set_ref=FeatureSetRef.from_feature_set(
FeatureSet.from_proto(feature_set_proto)
)
)
assert len(ingest_jobs) >= 1
ingest_job = ingest_jobs[0]
assert (
ingest_job.status == IngestionJobStatus.RUNNING
and ingest_job.id == "kafka-to-redis"
and ingest_job.external_id == "job-2222"
and ingest_job.feature_sets[0].name == "driver"
and ingest_job.source.source_type == "Kafka"
)
def test_get_feature_set(self, mock_client, mocker):
mock_client._core_service_stub = Core.CoreServiceStub(grpc.insecure_channel(""))
from google.protobuf.duration_pb2 import Duration
mocker.patch.object(
mock_client._core_service_stub,
"GetFeatureSet",
return_value=GetFeatureSetResponse(
feature_set=FeatureSetSpec(
name="my_feature_set",
version=2,
max_age=Duration(seconds=3600),
features=[
FeatureSpec(
name="my_feature_1", value_type=ValueProto.ValueType.FLOAT
),
FeatureSpec(
name="my_feature_2", value_type=ValueProto.ValueType.FLOAT
),
],
entities=[
EntitySpec(
name="my_entity_1", value_type=ValueProto.ValueType.INT64
)
],
source=Source(
type=SourceType.KAFKA,
kafka_source_config=KafkaSourceConfig(
bootstrap_servers="localhost:9092", topic="topic"
),
),
)
),
)
feature_set = mock_client.get_feature_set("my_feature_set", version=2)
assert (
feature_set.name == "my_feature_set"
and feature_set.version == 2
and feature_set.fields["my_feature_1"].name == "my_feature_1"
and feature_set.fields["my_feature_1"].dtype == ValueType.FLOAT
and feature_set.fields["my_entity_1"].name == "my_entity_1"
and feature_set.fields["my_entity_1"].dtype == ValueType.INT64
and len(feature_set.features) == 2
and len(feature_set.entities) == 1
)
def test_list_ingest_jobs(self, mock_jobcontroller_client, mocker):
mocker.patch.object(
mock_jobcontroller_client,
"_jobcontroller_service_stub",
return_value=Core.JobControllerServiceStub(grpc.insecure_channel("")),
)
feature_set_ref = FeatureSetRef(project="test", name="driver",)
mocker.patch.object(
mock_jobcontroller_client._jobcontroller_service_stub,
"ListIngestionJobs",
return_value=ListIngestionJobsResponse(
jobs=[
IngestJobProto(
id="kafka-to-redis",
external_id="job-2222",
status=IngestionJobStatus.RUNNING,
feature_set_references=[feature_set_ref.to_proto()],
source=Source(
type=SourceType.KAFKA,
kafka_source_config=KafkaSourceConfig(
bootstrap_servers="localhost:9092", topic="topic"
),
),
stores=[Store(name="redis")],
)
]
),
)
# list ingestion jobs by target feature set reference
ingest_jobs = mock_jobcontroller_client.list_ingest_jobs(
feature_set_ref=feature_set_ref
)
assert len(ingest_jobs) >= 1
ingest_job = ingest_jobs[0]
assert (
ingest_job.status == IngestionJobStatus.RUNNING
and ingest_job.id == "kafka-to-redis"
and ingest_job.external_id == "job-2222"
and ingest_job.feature_sets[0].name == "driver"
and ingest_job.source.source_type == "Kafka"
)
def test_stop_ingest_job(self, mocked_client, mocker):
mocker.patch.object(
mocked_client,
"_core_service_stub",
return_value=Core.CoreServiceStub(grpc.insecure_channel("")),
)
ingest_job = IngestJob(
job_proto=IngestJobProto(
id="kafka-to-redis",
external_id="job#2222",
status=IngestionJobStatus.RUNNING,
),
core_stub=mocked_client._core_service_stub,
)
mocked_client.stop_ingest_job(ingest_job)
assert mocked_client._core_service_stub.StopIngestionJob.called
def test_restart_ingest_job(self, mock_jobcontroller_client, mocker):
mocker.patch.object(
mock_jobcontroller_client,
"_jobcontroller_service_stub",
return_value=Core.JobControllerServiceStub(
grpc.insecure_channel("")),
)
ingest_job = IngestJob(
job_proto=IngestJobProto(
id="kafka-to-redis",
external_id="job#2222",
status=IngestionJobStatus.ERROR,
),
core_stub=mock_jobcontroller_client._jobcontroller_service_stub,
)
mock_jobcontroller_client.restart_ingest_job(ingest_job)
assert (mock_jobcontroller_client._jobcontroller_service_stub.
RestartIngestionJob.called)
def test_list_features(self, mocked_client, mocker):
mocker.patch.object(
mocked_client,
"_core_service_stub",
return_value=Core.CoreServiceStub(grpc.insecure_channel("")),
)
feature1_proto = FeatureSpecProto(
name="feature_1", value_type=ValueProto.ValueType.FLOAT
)
feature2_proto = FeatureSpecProto(
name="feature_2", value_type=ValueProto.ValueType.STRING
)
mocker.patch.object(
mocked_client._core_service_stub,
"ListFeatures",
return_value=ListFeaturesResponse(
features={
"driver_car:feature_1": feature1_proto,
"driver_car:feature_2": feature2_proto,
}
),
)
features = mocked_client.list_features_by_ref(project="test")
assert len(features) == 2
ref_str_list = []
feature_name_list = []
feature_dtype_list = []
for ref_str, feature_proto in features.items():
ref_str_list.append(ref_str)
feature_name_list.append(feature_proto.name)
feature_dtype_list.append(feature_proto.dtype)
assert (
set(ref_str_list) == set(["driver_car:feature_1", "driver_car:feature_2"])
and set(feature_name_list) == set(["feature_1", "feature_2"])
and set(feature_dtype_list) == set([ValueType.FLOAT, ValueType.STRING])
)
def test_get_batch_features(self, mock_client, mocker):
mock_client._serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel("")
)
mock_client._core_service_stub = Core.CoreServiceStub(grpc.insecure_channel(""))
mocker.patch.object(
mock_client._core_service_stub,
"GetFeatureSet",
return_value=GetFeatureSetResponse(
feature_set=FeatureSetSpec(
name="customer_fs",
version=1,
entities=[
EntitySpec(
name="customer", value_type=ValueProto.ValueType.INT64
),
EntitySpec(
name="transaction", value_type=ValueProto.ValueType.INT64
),
],
features=[
FeatureSpec(
name="customer_feature_1",
value_type=ValueProto.ValueType.FLOAT,
),
FeatureSpec(
name="customer_feature_2",
value_type=ValueProto.ValueType.STRING,
),
],
)
),
)
expected_dataframe = pd.DataFrame(
{
"datetime": [datetime.utcnow() for _ in range(3)],
"customer": [1001, 1002, 1003],
"transaction": [1001, 1002, 1003],
"customer_fs:1:customer_feature_1": [1001, 1002, 1003],
"customer_fs:1:customer_feature_2": [1001, 1002, 1003],
}
)
final_results = tempfile.mktemp()
to_avro(file_path_or_buffer=final_results, df=expected_dataframe)
mocker.patch.object(
mock_client._serving_service_stub,
"GetBatchFeatures",
return_value=GetBatchFeaturesResponse(
job=BatchFeaturesJob(
id="123",
type=JobType.JOB_TYPE_DOWNLOAD,
status=JobStatus.JOB_STATUS_DONE,
file_uris=[f"file://{final_results}"],
data_format=DataFormat.DATA_FORMAT_AVRO,
)
),
)
mocker.patch.object(
mock_client._serving_service_stub,
"GetJob",
return_value=GetJobResponse(
job=BatchFeaturesJob(
id="123",
type=JobType.JOB_TYPE_DOWNLOAD,
status=JobStatus.JOB_STATUS_DONE,
file_uris=[f"file://{final_results}"],
data_format=DataFormat.DATA_FORMAT_AVRO,
)
),
)
mocker.patch.object(
mock_client._serving_service_stub,
"GetFeastServingInfo",
return_value=GetFeastServingInfoResponse(
job_staging_location=f"file://{tempfile.mkdtemp()}/",
type=FeastServingType.FEAST_SERVING_TYPE_BATCH,
),
)
response = mock_client.get_batch_features(
entity_rows=pd.DataFrame(
{
"datetime": [
pd.datetime.now(tz=timezone("Asia/Singapore")) for _ in range(3)
],
"customer": [1001, 1002, 1003],
"transaction": [1001, 1002, 1003],
}
),
feature_ids=[
"customer_fs:1:customer_feature_1",
"customer_fs:1:customer_feature_2",
],
) # type: Job
assert response.id == "123" and response.status == JobStatus.JOB_STATUS_DONE
actual_dataframe = response.to_dataframe()
assert actual_dataframe[
["customer_fs:1:customer_feature_1", "customer_fs:1:customer_feature_2"]
].equals(
expected_dataframe[
["customer_fs:1:customer_feature_1", "customer_fs:1:customer_feature_2"]
]
)
def test_get_batch_features(self, mocked_client, mocker):
mocked_client._serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel(""))
mocked_client._core_service_stub = Core.CoreServiceStub(
grpc.insecure_channel(""))
mocker.patch.object(
mocked_client._core_service_stub,
"GetFeatureSet",
return_value=GetFeatureSetResponse(feature_set=FeatureSetProto(
spec=FeatureSetSpecProto(
name="driver",
project="driver_project",
entities=[
EntitySpecProto(name="driver",
value_type=ValueProto.ValueType.INT64),
EntitySpecProto(
name="transaction",
value_type=ValueProto.ValueType.INT64,
),
],
features=[
FeatureSpecProto(
name="driver_id",
value_type=ValueProto.ValueType.FLOAT,
),
FeatureSpecProto(
name="driver_name",
value_type=ValueProto.ValueType.STRING,
),
],
),
meta=FeatureSetMetaProto(
status=FeatureSetStatusProto.STATUS_READY),
)),
)
expected_dataframe = pd.DataFrame({
"datetime": [datetime.utcnow() for _ in range(3)],
"driver": [1001, 1002, 1003],
"transaction": [1001, 1002, 1003],
"driver_id": [1001, 1002, 1003],
})
final_results = tempfile.mktemp()
pandavro.to_avro(file_path_or_buffer=final_results,
df=expected_dataframe)
mocker.patch.object(
mocked_client._serving_service_stub,
"GetBatchFeatures",
return_value=GetBatchFeaturesResponse(job=BatchRetrievalJob(
id="123",
type=JobType.JOB_TYPE_DOWNLOAD,
status=JobStatus.JOB_STATUS_DONE,
file_uris=[f"file://{final_results}"],
data_format=DataFormat.DATA_FORMAT_AVRO,
)),
)
mocker.patch.object(
mocked_client._serving_service_stub,
"GetJob",
return_value=GetJobResponse(job=BatchRetrievalJob(
id="123",
type=JobType.JOB_TYPE_DOWNLOAD,
status=JobStatus.JOB_STATUS_DONE,
file_uris=[f"file://{final_results}"],
data_format=DataFormat.DATA_FORMAT_AVRO,
)),
)
mocker.patch.object(
mocked_client._serving_service_stub,
"GetFeastServingInfo",
return_value=GetFeastServingInfoResponse(
job_staging_location=f"file://{tempfile.mkdtemp()}/",
type=FeastServingType.FEAST_SERVING_TYPE_BATCH,
),
)
mocked_client.set_project("project1")
# TODO: Abstract away GCS client and GCP dependency
# NOTE: Feast Serving does not allow for feature references
# that specify the same feature in the same request.
with patch("google.cloud.storage.Client"):
response = mocked_client.get_batch_features(
entity_rows=pd.DataFrame({
"datetime": [
pd.datetime.now(tz=timezone("Asia/Singapore"))
for _ in range(3)
],
"driver": [1001, 1002, 1003],
"transaction": [1001, 1002, 1003],
}),
feature_refs=["driver:driver_id", "driver_id"],
project="driver_project",
) # Type: GetBatchFeaturesResponse
assert response.id == "123" and response.status == JobStatus.JOB_STATUS_DONE
actual_dataframe = response.to_dataframe()
assert actual_dataframe[["driver_id"
]].equals(expected_dataframe[["driver_id"]])
