def _entity_row_to_field_values(
row: GetOnlineFeaturesRequestV2.EntityRow,
) -> GetOnlineFeaturesResponse.FieldValues:
result = GetOnlineFeaturesResponse.FieldValues()
for k in row.fields:
result.fields[k].CopyFrom(row.fields[k])
result.statuses[k] = GetOnlineFeaturesResponse.FieldStatus.PRESENT
return result
def test_get_online_features(self, mocked_client, auth_metadata, mocker,
get_online_features_fields_statuses):
ROW_COUNT = 100
mocked_client._serving_service_stub = Serving.ServingServiceStub(
grpc.insecure_channel(""))
request = GetOnlineFeaturesRequestV2(project="driver_project")
request.features.extend([
FeatureRefProto(feature_table="driver", name="age"),
FeatureRefProto(feature_table="driver", name="rating"),
FeatureRefProto(feature_table="driver", name="null_value"),
])
receive_response = GetOnlineFeaturesResponse()
entity_rows = []
for row_number in range(0, ROW_COUNT):
fields = get_online_features_fields_statuses[row_number][0]
statuses = get_online_features_fields_statuses[row_number][1]
request.entity_rows.append(
GetOnlineFeaturesRequestV2.EntityRow(
fields={
"driver_id": ValueProto.Value(int64_val=row_number)
}))
entity_rows.append(
{"driver_id": ValueProto.Value(int64_val=row_number)})
receive_response.field_values.append(
GetOnlineFeaturesResponse.FieldValues(fields=fields,
statuses=statuses))
mocker.patch.object(
mocked_client._serving_service_stub,
"GetOnlineFeaturesV2",
return_value=receive_response,
)
got_response = mocked_client.get_online_features(
entity_rows=entity_rows,
feature_refs=["driver:age", "driver:rating", "driver:null_value"],
project="driver_project",
) # type: GetOnlineFeaturesResponse
mocked_client._serving_service_stub.GetOnlineFeaturesV2.assert_called_with(
request, metadata=auth_metadata, timeout=10)
got_fields = got_response.field_values[1].fields
got_statuses = got_response.field_values[1].statuses
assert (got_fields["driver_id"] == ValueProto.Value(int64_val=1)
and got_statuses["driver_id"]
== GetOnlineFeaturesResponse.FieldStatus.PRESENT
and got_fields["driver:age"] == ValueProto.Value(int64_val=1)
and got_statuses["driver:age"]
== GetOnlineFeaturesResponse.FieldStatus.PRESENT and
got_fields["driver:rating"] == ValueProto.Value(string_val="9")
and got_statuses["driver:rating"]
== GetOnlineFeaturesResponse.FieldStatus.PRESENT
and got_fields["driver:null_value"] == ValueProto.Value()
and got_statuses["driver:null_value"]
== GetOnlineFeaturesResponse.FieldStatus.NULL_VALUE)
