Skip to content

Commit f31d7f3

Browse files
terryyylimterryyylim
committed
Update tests
Signed-off-by: Terence Lim <terencelimxp@gmail.com>
1 parent 75d9369 commit f31d7f3

File tree

1 file changed

+83
-44
lines changed

1 file changed

+83
-44
lines changed

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

Lines changed: 83 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
import java.util.HashMap;
4242
import java.util.Map;
4343
import java.util.stream.Collectors;
44+
import java.util.stream.IntStream;
4445
import org.apache.avro.Schema;
4546
import org.apache.avro.SchemaBuilder;
4647
import org.apache.avro.generic.GenericDatumWriter;
@@ -217,48 +218,24 @@ static void globalSetup() throws IOException {
217218
.endRecord();
218219
byte[] schemaReference =
219220
Hashing.murmur3_32().hashBytes(ftSchema.toString().getBytes()).asBytes();
220-
221-
GenericRecord record =
222-
new GenericRecordBuilder(ftSchema)
223-
.set("trip_cost", 5)
224-
.set("trip_distance", 3.5)
225-
.set("trip_empty", null)
226-
.set("trip_wrong_type", "test")
227-
.build();
228-
byte[] entityFeatureKey =
229-
String.valueOf(DataGenerator.createInt64Value(1).getInt64Val()).getBytes();
230-
byte[] entityFeatureValue = createEntityValue(ftSchema, record);
231221
byte[] schemaKey = createSchemaKey(schemaReference);
232222

233-
ingestData(
234-
ridesFeatureTableName, cassandraTableName, entityFeatureKey, entityFeatureValue, schemaKey);
223+
ingestBulk(ridesFeatureTableName, cassandraTableName, ftSchema, 20);
235224

236225
Schema foodFtSchema =
237226
SchemaBuilder.record("FoodDriverData")
238227
.namespace(foodFeatureTableName)
239228
.fields()
240229
.requiredInt(feature1Reference.getName())
241230
.requiredDouble(feature2Reference.getName())
231+
.nullableString(feature3Reference.getName(), "null")
232+
.requiredString(feature4Reference.getName())
242233
.endRecord();
243234
byte[] foodSchemaReference =
244235
Hashing.murmur3_32().hashBytes(foodFtSchema.toString().getBytes()).asBytes();
245-
246-
GenericRecord foodRecord =
247-
new GenericRecordBuilder(foodFtSchema)
248-
.set("trip_cost", 12)
249-
.set("trip_distance", 7.5)
250-
.build();
251-
byte[] foodEntityFeatureKey =
252-
String.valueOf(DataGenerator.createInt64Value(1).getInt64Val()).getBytes();
253-
byte[] foodEntityFeatureValue = createEntityValue(foodFtSchema, foodRecord);
254236
byte[] foodSchemaKey = createSchemaKey(foodSchemaReference);
255237

256-
ingestData(
257-
foodFeatureTableName,
258-
cassandraTableName,
259-
foodEntityFeatureKey,
260-
foodEntityFeatureValue,
261-
foodSchemaKey);
238+
ingestBulk(foodFeatureTableName, cassandraTableName, foodFtSchema, 20);
262239

263240
/** Compound Entity Ingestion Workflow */
264241
Schema compoundFtSchema =
@@ -376,6 +353,40 @@ private static void ingestData(
376353
ByteBuffer.wrap(entityFeatureValue)));
377354
}
378355

356+
private static void ingestBulk(
357+
String featureTableName, String cassandraTableName, Schema schema, Integer counts) {
358+
359+
IntStream.range(0, counts)
360+
.forEach(
361+
i -> {
362+
try {
363+
GenericRecord record =
364+
new GenericRecordBuilder(schema)
365+
.set("trip_cost", i)
366+
.set("trip_distance", (double) i)
367+
.set("trip_empty", null)
368+
.set("trip_wrong_type", "test")
369+
.build();
370+
byte[] schemaReference =
371+
Hashing.murmur3_32().hashBytes(schema.toString().getBytes()).asBytes();
372+
373+
byte[] entityFeatureKey =
374+
String.valueOf(DataGenerator.createInt64Value(i).getInt64Val()).getBytes();
375+
byte[] entityFeatureValue = createEntityValue(schema, record);
376+
377+
byte[] schemaKey = createSchemaKey(schemaReference);
378+
ingestData(
379+
featureTableName,
380+
cassandraTableName,
381+
entityFeatureKey,
382+
entityFeatureValue,
383+
schemaKey);
384+
} catch (IOException e) {
385+
e.printStackTrace();
386+
}
387+
});
388+
}
389+
379390
private static void ingestSchema(byte[] schemaKey, Schema schema) {
380391
PreparedStatement schemaStatement =
381392
cqlSession.prepare(
@@ -433,7 +444,7 @@ public void shouldRegisterSingleEntityAndGetOnlineFeatures() {
433444
entityName,
434445
entityValue,
435446
FeatureV2.getFeatureStringRef(featureReference),
436-
DataGenerator.createInt32Value(5),
447+
DataGenerator.createInt32Value(1),
437448
FeatureV2.getFeatureStringRef(notFoundFeatureReference),
438449
DataGenerator.createEmptyValue());
439450

@@ -496,7 +507,7 @@ public void shouldRegisterCompoundEntityAndGetOnlineFeatures() {
496507
merchantEntityName,
497508
merchantEntityValue,
498509
FeatureV2.getFeatureStringRef(featureReference),
499-
DataGenerator.createInt32Value(5),
510+
DataGenerator.createInt32Value(1),
500511
FeatureV2.getFeatureStringRef(notFoundFeatureReference),
501512
DataGenerator.createEmptyValue());
502513

@@ -523,20 +534,26 @@ public void shouldRegisterCompoundEntityAndGetOnlineFeatures() {
523534
}
524535

525536
@Test
526-
public void shouldReturnCorrectRowCount() {
537+
public void shouldReturnCorrectRowCountAndOrder() {
527538
// getOnlineFeatures Information
528539
String projectName = "default";
529540
String entityName = "driver_id";
530541
ValueProto.Value entityValue1 = ValueProto.Value.newBuilder().setInt64Val(1).build();
531542
ValueProto.Value entityValue2 = ValueProto.Value.newBuilder().setInt64Val(2).build();
543+
ValueProto.Value entityValue3 = ValueProto.Value.newBuilder().setInt64Val(3).build();
544+
ValueProto.Value entityValue4 = ValueProto.Value.newBuilder().setInt64Val(4).build();
532545

533546
// Instantiate EntityRows
534547
GetOnlineFeaturesRequestV2.EntityRow entityRow1 =
535548
DataGenerator.createEntityRow(entityName, entityValue1, 100);
536549
GetOnlineFeaturesRequestV2.EntityRow entityRow2 =
537550
DataGenerator.createEntityRow(entityName, entityValue2, 100);
551+
GetOnlineFeaturesRequestV2.EntityRow entityRow3 =
552+
DataGenerator.createEntityRow(entityName, entityValue3, 100);
553+
GetOnlineFeaturesRequestV2.EntityRow entityRow4 =
554+
DataGenerator.createEntityRow(entityName, entityValue4, 100);
538555
ImmutableList<GetOnlineFeaturesRequestV2.EntityRow> entityRows =
539-
ImmutableList.of(entityRow1, entityRow2);
556+
ImmutableList.of(entityRow1, entityRow2, entityRow4, entityRow3);
540557

541558
// Instantiate FeatureReferences
542559
FeatureReferenceV2 featureReference =
@@ -560,7 +577,7 @@ public void shouldReturnCorrectRowCount() {
560577
entityName,
561578
entityValue1,
562579
FeatureV2.getFeatureStringRef(featureReference),
563-
DataGenerator.createInt32Value(5),
580+
DataGenerator.createInt32Value(1),
564581
FeatureV2.getFeatureStringRef(notFoundFeatureReference),
565582
DataGenerator.createEmptyValue(),
566583
FeatureV2.getFeatureStringRef(emptyFeatureReference),
@@ -588,30 +605,52 @@ public void shouldReturnCorrectRowCount() {
588605
entityName,
589606
entityValue2,
590607
FeatureV2.getFeatureStringRef(featureReference),
608+
DataGenerator.createInt32Value(2),
609+
FeatureV2.getFeatureStringRef(notFoundFeatureReference),
591610
DataGenerator.createEmptyValue(),
611+
FeatureV2.getFeatureStringRef(emptyFeatureReference),
612+
DataGenerator.createEmptyValue());
613+
614+
ImmutableMap<String, ValueProto.Value> expectedValueMap3 =
615+
ImmutableMap.of(
616+
entityName,
617+
entityValue3,
618+
FeatureV2.getFeatureStringRef(featureReference),
619+
DataGenerator.createInt32Value(3),
592620
FeatureV2.getFeatureStringRef(notFoundFeatureReference),
593621
DataGenerator.createEmptyValue(),
594622
FeatureV2.getFeatureStringRef(emptyFeatureReference),
595623
DataGenerator.createEmptyValue());
596624

597-
ImmutableMap<String, GetOnlineFeaturesResponse.FieldStatus> expectedStatusMap2 =
625+
ImmutableMap<String, ValueProto.Value> expectedValueMap4 =
598626
ImmutableMap.of(
599627
entityName,
600-
GetOnlineFeaturesResponse.FieldStatus.PRESENT,
628+
entityValue4,
601629
FeatureV2.getFeatureStringRef(featureReference),
602-
GetOnlineFeaturesResponse.FieldStatus.NOT_FOUND,
630+
DataGenerator.createInt32Value(4),
603631
FeatureV2.getFeatureStringRef(notFoundFeatureReference),
604-
GetOnlineFeaturesResponse.FieldStatus.NOT_FOUND,
632+
DataGenerator.createEmptyValue(),
605633
FeatureV2.getFeatureStringRef(emptyFeatureReference),
606-
GetOnlineFeaturesResponse.FieldStatus.NOT_FOUND);
634+
DataGenerator.createEmptyValue());
607635

608636
GetOnlineFeaturesResponse.FieldValues expectedFieldValues2 =
609637
GetOnlineFeaturesResponse.FieldValues.newBuilder()
610638
.putAllFields(expectedValueMap2)
611-
.putAllStatuses(expectedStatusMap2)
639+
.putAllStatuses(expectedStatusMap)
640+
.build();
641+
GetOnlineFeaturesResponse.FieldValues expectedFieldValues3 =
642+
GetOnlineFeaturesResponse.FieldValues.newBuilder()
643+
.putAllFields(expectedValueMap3)
644+
.putAllStatuses(expectedStatusMap)
645+
.build();
646+
GetOnlineFeaturesResponse.FieldValues expectedFieldValues4 =
647+
GetOnlineFeaturesResponse.FieldValues.newBuilder()
648+
.putAllFields(expectedValueMap4)
649+
.putAllStatuses(expectedStatusMap)
612650
.build();
613651
ImmutableList<GetOnlineFeaturesResponse.FieldValues> expectedFieldValuesList =
614-
ImmutableList.of(expectedFieldValues, expectedFieldValues2);
652+
ImmutableList.of(
653+
expectedFieldValues, expectedFieldValues2, expectedFieldValues4, expectedFieldValues3);
615654

616655
assertEquals(expectedFieldValuesList, featureResponse.getFieldValuesList());
617656
}
@@ -655,13 +694,13 @@ public void shouldReturnFeaturesFromDiffFeatureTable() {
655694
entityName,
656695
entityValue,
657696
FeatureV2.getFeatureStringRef(rideFeatureReference),
658-
DataGenerator.createInt32Value(5),
697+
DataGenerator.createInt32Value(1),
659698
FeatureV2.getFeatureStringRef(rideFeatureReference2),
660-
DataGenerator.createDoubleValue(3.5),
699+
DataGenerator.createDoubleValue(1.0),
661700
FeatureV2.getFeatureStringRef(foodFeatureReference),
662-
DataGenerator.createInt32Value(12),
701+
DataGenerator.createInt32Value(1),
663702
FeatureV2.getFeatureStringRef(foodFeatureReference2),
664-
DataGenerator.createDoubleValue(7.5));
703+
DataGenerator.createDoubleValue(1.0));
665704

666705
ImmutableMap<String, GetOnlineFeaturesResponse.FieldStatus> expectedStatusMap =
667706
ImmutableMap.of(

0 commit comments

Comments
 (0)