Skip to content

Commit 5d765cb

Browse files
davidheryantodavidheryanto
committed
Add integration test for import job with Redis store.
Given a valid FeatureSetSpec and FeatureRow elements, make sure the same FeatureRow for the corresponding FeatureSetSpec can be retrieved. The integration test is accomplished by running embedded Kafka and Redis, then running Beam pipeline with DirectRunner. This commit also adds TestUtil class that helps with setting up local Kafka, Redis and creating test data.
1 parent f985024 commit 5d765cb

5 files changed

Lines changed: 445 additions & 2 deletions

File tree

‎ingestion/pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -299,11 +299,11 @@
299299
<scope>test</scope>
300300
</dependency>
301301

302+
<!-- To simulate actual Redis for ingestion integration test -->
302303
<dependency>
303304
<groupId>com.github.kstyrc</groupId>
304305
<artifactId>embedded-redis</artifactId>
305306
<version>0.6</version>
306-
<scope>test</scope>
307307
</dependency>
308308

309309
<dependency>

‎ingestion/src/main/java/feast/store/serving/redis/FeatureRowToRedisMutationDoFn.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ private RedisKey getKey(FeatureRow featureRow) {
5050
redisKeyBuilder.addEntities(field);
5151
}
5252
}
53+
5354
return redisKeyBuilder.build();
5455
}
5556

‎ingestion/src/test/java/feast/ingestion/ImportJobTest.java‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
package feast.ingestion;
2+
3+
import com.google.common.io.Files;
4+
import com.google.protobuf.InvalidProtocolBufferException;
5+
import com.google.protobuf.util.JsonFormat;
6+
import feast.core.FeatureSetProto.EntitySpec;
7+
import feast.core.FeatureSetProto.FeatureSetSpec;
8+
import feast.core.FeatureSetProto.FeatureSpec;
9+
import feast.core.SourceProto.KafkaSourceConfig;
10+
import feast.core.SourceProto.Source;
11+
import feast.core.SourceProto.SourceType;
12+
import feast.core.StoreProto.Store;
13+
import feast.core.StoreProto.Store.RedisConfig;
14+
import feast.core.StoreProto.Store.StoreType;
15+
import feast.core.StoreProto.Store.Subscription;
16+
import feast.ingestion.options.ImportOptions;
17+
import feast.storage.RedisProto.RedisKey;
18+
import feast.test.TestUtil;
19+
import feast.test.TestUtil.LocalKafka;
20+
import feast.test.TestUtil.LocalRedis;
21+
import feast.types.FeatureRowProto.FeatureRow;
22+
import feast.types.ValueProto.ValueType.Enum;
23+
import java.io.IOException;
24+
import java.nio.charset.StandardCharsets;
25+
import java.time.Duration;
26+
import java.util.ArrayList;
27+
import java.util.Collections;
28+
import java.util.HashMap;
29+
import java.util.List;
30+
import java.util.Map;
31+
import java.util.stream.IntStream;
32+
import org.apache.beam.sdk.PipelineResult;
33+
import org.apache.beam.sdk.PipelineResult.State;
34+
import org.apache.beam.sdk.options.PipelineOptionsFactory;
35+
import org.apache.kafka.common.serialization.ByteArraySerializer;
36+
import org.junit.AfterClass;
37+
import org.junit.Assert;
38+
import org.junit.BeforeClass;
39+
import org.junit.Test;
40+
import org.slf4j.Logger;
41+
import org.slf4j.LoggerFactory;
42+
import redis.clients.jedis.Jedis;
43+
44+
public class ImportJobTest {
45+
46+
private static final Logger LOGGER = LoggerFactory.getLogger(ImportJobTest.class.getName());
47+
48+
private static final String KAFKA_HOST = "localhost";
49+
private static final int KAFKA_PORT = 9093;
50+
private static final String KAFKA_BOOTSTRAP_SERVERS = KAFKA_HOST + ":" + KAFKA_PORT;
51+
private static final short KAFKA_REPLICATION_FACTOR = 1;
52+
private static final String KAFKA_TOPIC = "topic_1";
53+
private static final long KAFKA_PUBLISH_TIMEOUT_SEC = 10;
54+
55+
@SuppressWarnings("UnstableApiUsage")
56+
private static final String ZOOKEEPER_DATA_DIR = Files.createTempDir().getAbsolutePath();
57+
private static final String ZOOKEEPER_HOST = "localhost";
58+
private static final int ZOOKEEPER_PORT = 2182;
59+
60+
private static final String REDIS_HOST = "localhost";
61+
private static final int REDIS_PORT = 6380;
62+
63+
// Expected time taken for the import job to be ready to receive Feature Row input
64+
private static final int IMPORT_JOB_READY_DURATION_SEC = 5;
65+
// Expected time taken for the import job to finish writing to Store
66+
private static final int IMPORT_JOB_RUN_DURATION_SEC = 10;
67+
68+
@BeforeClass
69+
public static void setup() throws IOException, InterruptedException {
70+
LocalKafka.start(KAFKA_HOST, KAFKA_PORT, KAFKA_REPLICATION_FACTOR, true, ZOOKEEPER_HOST,
71+
ZOOKEEPER_PORT, ZOOKEEPER_DATA_DIR);
72+
LocalRedis.start(REDIS_PORT);
73+
}
74+
75+
@AfterClass
76+
public static void tearDown() {
77+
LocalRedis.stop();
78+
LocalKafka.stop();
79+
}
80+
81+
@Test
82+
public void runPipeline_ShouldWriteToRedisCorrectlyGivenValidSpecAndFeatureRow()
83+
throws IOException, InterruptedException {
84+
FeatureSetSpec spec =
85+
FeatureSetSpec.newBuilder().setName("feature_set").setVersion(3)
86+
.addEntities(EntitySpec.newBuilder()
87+
.setName("entity_id_primary").setValueType(Enum.INT32).build())
88+
.addEntities(EntitySpec.newBuilder()
89+
.setName("entity_id_secondary").setValueType(Enum.STRING).build())
90+
.addFeatures(FeatureSpec.newBuilder()
91+
.setName("feature_1").setValueType(Enum.STRING_LIST).build())
92+
.addFeatures(FeatureSpec.newBuilder()
93+
.setName("feature_2").setValueType(Enum.STRING).build())
94+
.addFeatures(FeatureSpec.newBuilder()
95+
.setName("feature_3").setValueType(Enum.INT64).build())
96+
.setSource(Source.newBuilder()
97+
.setType(SourceType.KAFKA).setKafkaSourceConfig(
98+
KafkaSourceConfig.newBuilder()
99+
.setBootstrapServers(KAFKA_HOST + ":" + KAFKA_PORT)
100+
.setTopic(KAFKA_TOPIC).build())
101+
.build())
102+
.build();
103+
104+
Store redis =
105+
Store.newBuilder().setName(StoreType.REDIS.toString()).setType(StoreType.REDIS)
106+
.setRedisConfig(RedisConfig.newBuilder()
107+
.setHost(REDIS_HOST).setPort(REDIS_PORT).build())
108+
.addSubscriptions(Subscription.newBuilder()
109+
.setName(spec.getName()).setVersion(String.valueOf(spec.getVersion())).build())
110+
.build();
111+
112+
ImportOptions options = PipelineOptionsFactory.create().as(ImportOptions.class);
113+
options.setFeatureSetSpecJson(
114+
Collections.singletonList(
115+
JsonFormat.printer().omittingInsignificantWhitespace().print(spec)));
116+
options.setStoreJson(
117+
Collections.singletonList(
118+
JsonFormat.printer().omittingInsignificantWhitespace().print(redis)));
119+
options.setBlockOnRun(false);
120+
121+
int inputSize = 4096;
122+
List<FeatureRow> input = new ArrayList<>();
123+
Map<RedisKey, FeatureRow> expected = new HashMap<>();
124+
125+
LOGGER.info("Generating test data ...");
126+
IntStream.range(0, inputSize).forEach(i -> {
127+
FeatureRow randomRow = TestUtil.createRandomFeatureRow(spec);
128+
RedisKey redisKey = TestUtil.createRedisKey(spec, randomRow);
129+
input.add(randomRow);
130+
expected.put(redisKey, randomRow);
131+
});
132+
133+
LOGGER.info("Starting Import Job with the following options: {}", options.toString());
134+
PipelineResult pipelineResult = ImportJob.runPipeline(options);
135+
Thread.sleep(Duration.ofSeconds(IMPORT_JOB_READY_DURATION_SEC).toMillis());
136+
Assert.assertEquals(pipelineResult.getState(), State.RUNNING);
137+
138+
LOGGER.info("Publishing {} Feature Row messages to Kafka ...", input.size());
139+
TestUtil.publishFeatureRowsToKafka(KAFKA_BOOTSTRAP_SERVERS, KAFKA_TOPIC, input,
140+
ByteArraySerializer.class, KAFKA_PUBLISH_TIMEOUT_SEC);
141+
Thread.sleep(Duration.ofSeconds(IMPORT_JOB_RUN_DURATION_SEC).toMillis());
142+
143+
LOGGER.info("Validating the actual values written to Redis ...");
144+
Jedis jedis = new Jedis(REDIS_HOST, REDIS_PORT);
145+
expected.forEach((key, expectedValue) -> {
146+
byte[] actualByteValue = jedis.get(key.toByteArray());
147+
Assert.assertNotNull("Key not found in Redis: " + key, actualByteValue);
148+
FeatureRow actualValue = null;
149+
try {
150+
actualValue = FeatureRow.parseFrom(actualByteValue);
151+
} catch (InvalidProtocolBufferException e) {
152+
Assert.fail(String
153+
.format("Actual Redis value cannot be parsed as FeatureRow, key: %s, value :%s",
154+
key, new String(actualByteValue, StandardCharsets.UTF_8)));
155+
}
156+
Assert.assertEquals(expectedValue, actualValue);
157+
});
158+
}
159+
}

0 commit comments

Comments
 (0)