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)