|
25 | 25 | import static org.mockito.Mockito.when;
|
26 | 26 | import static org.mockito.MockitoAnnotations.initMocks;
|
27 | 27 |
|
| 28 | +import com.google.common.collect.Lists; |
28 | 29 | import com.google.protobuf.InvalidProtocolBufferException;
|
29 | 30 | import feast.core.CoreServiceProto.ListFeatureSetsRequest.Filter;
|
30 | 31 | import feast.core.CoreServiceProto.ListFeatureSetsResponse;
|
@@ -194,6 +195,13 @@ public void shouldGenerateAndSubmitJobsIfAny() throws InvalidProtocolBufferExcep
|
194 | 195 | when(specService.listStores(any()))
|
195 | 196 | .thenReturn(ListStoresResponse.newBuilder().addStore(store).build());
|
196 | 197 |
|
| 198 | + for (FeatureSetProto.FeatureSet fs : Lists.newArrayList(featureSet1, featureSet2)) { |
| 199 | + FeatureSetSpec spec = fs.getSpec(); |
| 200 | + when(featureSetRepository.findFeatureSetByNameAndProject_NameAndVersion( |
| 201 | + spec.getName(), spec.getProject(), spec.getVersion())) |
| 202 | + .thenReturn(FeatureSet.fromProto(fs)); |
| 203 | + } |
| 204 | + |
197 | 205 | when(jobManager.startJob(argThat(new JobMatcher(expectedInput)))).thenReturn(expected);
|
198 | 206 | when(jobManager.getRunnerType()).thenReturn(Runner.DATAFLOW);
|
199 | 207 |
|
@@ -318,6 +326,12 @@ public void shouldGroupJobsBySource() throws InvalidProtocolBufferException {
|
318 | 326 | when(jobManager.startJob(argThat(new JobMatcher(expectedInput1)))).thenReturn(expected1);
|
319 | 327 | when(jobManager.startJob(argThat(new JobMatcher(expectedInput2)))).thenReturn(expected2);
|
320 | 328 | when(jobManager.getRunnerType()).thenReturn(Runner.DATAFLOW);
|
| 329 | + for (FeatureSetProto.FeatureSet fs : Lists.newArrayList(featureSet1, featureSet2)) { |
| 330 | + FeatureSetSpec spec = fs.getSpec(); |
| 331 | + when(featureSetRepository.findFeatureSetByNameAndProject_NameAndVersion( |
| 332 | + spec.getName(), spec.getProject(), spec.getVersion())) |
| 333 | + .thenReturn(FeatureSet.fromProto(fs)); |
| 334 | + } |
321 | 335 |
|
322 | 336 | JobCoordinatorService jcs =
|
323 | 337 | new JobCoordinatorService(
|
|
0 commit comments