Skip to content

Commit c278b91

Browse files
terryyylimterryyylim
committed
Add multi entityRows IT
Signed-off-by: Terence Lim <terencelimxp@gmail.com>
1 parent b6d8201 commit c278b91

1 file changed

Lines changed: 94 additions & 0 deletions

File tree

‎serving/src/test/java/feast/serving/it/ServingServiceBigTableIT.java‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -510,6 +510,100 @@ public void shouldRegisterCompoundEntityAndGetOnlineFeatures() {
510510
assertEquals(expectedFieldValuesList, featureResponse.getFieldValuesList());
511511
}
512512

513+
@Test
514+
public void shouldReturnCorrectRowCount() {
515+
// getOnlineFeatures Information
516+
String projectName = "default";
517+
String entityName = "driver_id";
518+
ValueProto.Value entityValue1 = ValueProto.Value.newBuilder().setInt64Val(1).build();
519+
ValueProto.Value entityValue2 = ValueProto.Value.newBuilder().setInt64Val(2).build();
520+
521+
// Instantiate EntityRows
522+
ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow entityRow1 =
523+
DataGenerator.createEntityRow(entityName, entityValue1, 100);
524+
ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow entityRow2 =
525+
DataGenerator.createEntityRow(entityName, entityValue2, 100);
526+
ImmutableList<ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow> entityRows =
527+
ImmutableList.of(entityRow1, entityRow2);
528+
529+
// Instantiate FeatureReferences
530+
ServingAPIProto.FeatureReferenceV2 featureReference =
531+
DataGenerator.createFeatureReference("rides", "trip_cost");
532+
ServingAPIProto.FeatureReferenceV2 notFoundFeatureReference =
533+
DataGenerator.createFeatureReference("rides", "trip_transaction");
534+
ServingAPIProto.FeatureReferenceV2 emptyFeatureReference =
535+
DataGenerator.createFeatureReference("rides", "trip_empty");
536+
537+
ImmutableList<ServingAPIProto.FeatureReferenceV2> featureReferences =
538+
ImmutableList.of(featureReference, notFoundFeatureReference, emptyFeatureReference);
539+
540+
// Build GetOnlineFeaturesRequestV2
541+
ServingAPIProto.GetOnlineFeaturesRequestV2 onlineFeatureRequest =
542+
TestUtils.createOnlineFeatureRequest(projectName, featureReferences, entityRows);
543+
ServingAPIProto.GetOnlineFeaturesResponse featureResponse =
544+
servingStub.getOnlineFeaturesV2(onlineFeatureRequest);
545+
546+
ImmutableMap<String, ValueProto.Value> expectedValueMap =
547+
ImmutableMap.of(
548+
entityName,
549+
entityValue1,
550+
FeatureV2.getFeatureStringRef(featureReference),
551+
DataGenerator.createInt64Value(5),
552+
FeatureV2.getFeatureStringRef(notFoundFeatureReference),
553+
DataGenerator.createEmptyValue(),
554+
FeatureV2.getFeatureStringRef(emptyFeatureReference),
555+
DataGenerator.createEmptyValue());
556+
557+
ImmutableMap<String, ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus> expectedStatusMap =
558+
ImmutableMap.of(
559+
entityName,
560+
ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus.PRESENT,
561+
FeatureV2.getFeatureStringRef(featureReference),
562+
ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus.PRESENT,
563+
FeatureV2.getFeatureStringRef(notFoundFeatureReference),
564+
ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus.NOT_FOUND,
565+
FeatureV2.getFeatureStringRef(emptyFeatureReference),
566+
ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus.NOT_FOUND);
567+
568+
ServingAPIProto.GetOnlineFeaturesResponse.FieldValues expectedFieldValues =
569+
ServingAPIProto.GetOnlineFeaturesResponse.FieldValues.newBuilder()
570+
.putAllFields(expectedValueMap)
571+
.putAllStatuses(expectedStatusMap)
572+
.build();
573+
574+
ImmutableMap<String, ValueProto.Value> expectedValueMap2 =
575+
ImmutableMap.of(
576+
entityName,
577+
entityValue2,
578+
FeatureV2.getFeatureStringRef(featureReference),
579+
DataGenerator.createEmptyValue(),
580+
FeatureV2.getFeatureStringRef(notFoundFeatureReference),
581+
DataGenerator.createEmptyValue(),
582+
FeatureV2.getFeatureStringRef(emptyFeatureReference),
583+
DataGenerator.createEmptyValue());
584+
585+
ImmutableMap<String, ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus> expectedStatusMap2 =
586+
ImmutableMap.of(
587+
entityName,
588+
ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus.PRESENT,
589+
FeatureV2.getFeatureStringRef(featureReference),
590+
ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus.NOT_FOUND,
591+
FeatureV2.getFeatureStringRef(notFoundFeatureReference),
592+
ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus.NOT_FOUND,
593+
FeatureV2.getFeatureStringRef(emptyFeatureReference),
594+
ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus.NOT_FOUND);
595+
596+
ServingAPIProto.GetOnlineFeaturesResponse.FieldValues expectedFieldValues2 =
597+
ServingAPIProto.GetOnlineFeaturesResponse.FieldValues.newBuilder()
598+
.putAllFields(expectedValueMap2)
599+
.putAllStatuses(expectedStatusMap2)
600+
.build();
601+
ImmutableList<ServingAPIProto.GetOnlineFeaturesResponse.FieldValues> expectedFieldValuesList =
602+
ImmutableList.of(expectedFieldValues, expectedFieldValues2);
603+
604+
assertEquals(expectedFieldValuesList, featureResponse.getFieldValuesList());
605+
}
606+
513607
@TestConfiguration
514608
public static class TestConfig {
515609
@Bean

0 commit comments

Comments
 (0)