@Test
public void testReadSingletonSideInputValue() throws Exception {
CloudObject sideInputKind = CloudObject.forClassName("singleton");
Object elem = "hi";
List<Object> elems = Arrays.asList(elem);
assertEquals(elem, SideInputUtils.readSideInputValue(sideInputKind, elems));
}
@Test
public void testCreateNormalParDoFn() throws Exception {
String stringState = "some state";
long longState = 42L;
TestDoFn fn = new TestDoFn(stringState, longState);
String serializedFn =
StringUtils.byteArrayToJsonString(
SerializableUtils.serializeToByteArray(
new DoFnInfo(fn, WindowingStrategy.globalDefault())));
CloudObject cloudUserFn = CloudObject.forClassName("DoFn");
addString(cloudUserFn, "serialized_fn", serializedFn);
String tag = "output";
MultiOutputInfo multiOutputInfo = new MultiOutputInfo();
multiOutputInfo.setTag(tag);
List<MultiOutputInfo> multiOutputInfos = Arrays.asList(multiOutputInfo);
PipelineOptions options = PipelineOptionsFactory.create();
DataflowExecutionContext context = BatchModeExecutionContext.fromOptions(options);
CounterSet counters = new CounterSet();
StateSampler stateSampler = new StateSampler("test", counters.getAddCounterMutator());
ParDoFn parDoFn =
factory.create(
options,
cloudUserFn,
"name",
"transformName",
null,
multiOutputInfos,
1,
context,
counters.getAddCounterMutator(),
stateSampler);
// Test that the factory created the correct class
assertThat(parDoFn, instanceOf(NormalParDoFn.class));
// Test that the DoFnInfo reflects the one passed in
NormalParDoFn normalParDoFn = (NormalParDoFn) parDoFn;
DoFnInfo doFnInfo = normalParDoFn.getDoFnInfo();
DoFn actualDoFn = doFnInfo.getDoFn();
assertThat(actualDoFn, instanceOf(TestDoFn.class));
assertThat(doFnInfo.getWindowingStrategy().getWindowFn(), instanceOf(GlobalWindows.class));
assertThat(
doFnInfo.getWindowingStrategy().getTrigger().getSpec(), instanceOf(DefaultTrigger.class));
// Test that the deserialized user DoFn is as expected
TestDoFn actualTestDoFn = (TestDoFn) actualDoFn;
assertEquals(stringState, actualTestDoFn.stringState);
assertEquals(longState, actualTestDoFn.longState);
assertEquals(context, normalParDoFn.getExecutionContext());
}
@Test
public void testCreateUnknownParDoFn() throws Exception {
CloudObject cloudUserFn = CloudObject.forClassName("UnknownKindOfDoFn");
try {
CounterSet counters = new CounterSet();
StateSampler stateSampler = new StateSampler("test", counters.getAddCounterMutator());
factory.create(
PipelineOptionsFactory.create(),
cloudUserFn,
"name",
"transformName",
null,
null,
1,
BatchModeExecutionContext.fromOptions(PipelineOptionsFactory.create()),
counters.getAddCounterMutator(),
stateSampler);
fail("should have thrown an exception");
} catch (Exception exn) {
assertThat(exn.toString(), Matchers.containsString("No known ParDoFnFactory"));
}
}
@SuppressWarnings("rawtypes")
private static ParDoFn createCombineValuesFn(String phase, Combine.KeyedCombineFn combineFn)
throws Exception {
// This partially mirrors the work that
// com.google.cloud.dataflow.sdk.transforms.Combine.translateHelper
// does, at least for the KeyedCombineFn. The phase is generated
// by the back-end.
CloudObject spec = CloudObject.forClassName("CombineValuesFn");
addString(
spec, PropertyNames.SERIALIZED_FN, byteArrayToJsonString(serializeToByteArray(combineFn)));
addString(spec, PropertyNames.PHASE, phase);
return parDoFnFactory.create(
PipelineOptionsFactory.create(),
spec,
"name",
"transformName",
null, // no side inputs
null, // no side outputs
1, // single main output
DataflowExecutionContext.withoutSideInputs(),
(new CounterSet()).getAddCounterMutator(),
null);
}
SideInputInfo createCollectionSideInputInfo(Source... sideInputSources) {
SideInputInfo sideInputInfo = new SideInputInfo();
sideInputInfo.setSources(Arrays.asList(sideInputSources));
sideInputInfo.setKind(CloudObject.forClassName("collection"));
return sideInputInfo;
}
SideInputInfo createSingletonSideInputInfo(Source sideInputSource) {
SideInputInfo sideInputInfo = new SideInputInfo();
sideInputInfo.setSources(Arrays.asList(sideInputSource));
sideInputInfo.setKind(CloudObject.forClassName("singleton"));
return sideInputInfo;
}
@Test
public void testReadCollectionSideInputValue() throws Exception {
CloudObject sideInputKind = CloudObject.forClassName("collection");
List<Object> elems = Arrays.<Object>asList("hi", "there", "bob");
assertEquals(elems, SideInputUtils.readSideInputValue(sideInputKind, elems));
}