FixMultiOutputInfosOnParDoInstructions implements Function<MapTask, MapTask> { @Override public MapTask apply(MapTask input) { for (ParallelInstruction instruction : Apiary.listOrEmpty(input.getInstructions())) { ParDoInstruction parDoInstruction = instruction.getParDo(); if (parDoInstruction != null) { int numOutputs = Apiary.intOrZero(parDoInstruction.getNumOutputs()); List<MultiOutputInfo> multiOutputInfos = Apiary.listOrEmpty(parDoInstruction.getMultiOutputInfos()); if (numOutputs != Apiary.listOrEmpty(instruction.getParDo().getMultiOutputInfos()).size()) { if (numOutputs == 1) { parDoInstruction.setMultiOutputInfos( ImmutableList.of(new MultiOutputInfo().setTag(idGenerator.getId()))); } else { throw new IllegalArgumentException( String.format( "Invalid ParDoInstruction %s, %d outputs specified, found %s tags.", instruction.getSystemName(), numOutputs, multiOutputInfos)); } } } } return input; } FixMultiOutputInfosOnParDoInstructions(IdGenerator idGenerator); @Override MapTask apply(MapTask input);  }
@Test public void testExistingMultiOutputInfosAreUnmodified() { FixMultiOutputInfosOnParDoInstructions function = new FixMultiOutputInfosOnParDoInstructions(IdGenerators.decrementingLongs()); MapTask output = function.apply(createMapTaskWithParDo(2, "5", "6")); assertEquals(createMapTaskWithParDo(2, "5", "6"), output); }
@Test public void testDefaultOutputIsAddedIfOnlySingleOutput() { FixMultiOutputInfosOnParDoInstructions function = new FixMultiOutputInfosOnParDoInstructions(IdGenerators.decrementingLongs()); MapTask output = function.apply(createMapTaskWithParDo(1)); assertEquals(createMapTaskWithParDo(1, "-1"), output); }
@Test public void testDefaultOutputHasDifferentIdsForEachMapTask() { FixMultiOutputInfosOnParDoInstructions function = new FixMultiOutputInfosOnParDoInstructions(IdGenerators.decrementingLongs()); MapTask output = function.apply(createMapTaskWithParDo(1)); assertEquals(createMapTaskWithParDo(1, "-1"), output); output = function.apply(createMapTaskWithParDo(1)); assertEquals(createMapTaskWithParDo(1, "-2"), output); }
@Test public void testMissingTagsForMultipleOutputsThrows() { FixMultiOutputInfosOnParDoInstructions function = new FixMultiOutputInfosOnParDoInstructions(IdGenerators.decrementingLongs()); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("Invalid ParDoInstruction"); thrown.expectMessage("2 outputs specified"); function.apply(createMapTaskWithParDo(2)); }