LengthPrefixUnknownCoders { @VisibleForTesting static ParallelInstruction forParallelInstruction( ParallelInstruction input, boolean replaceWithByteArrayCoder) throws Exception { try { ParallelInstruction instruction = clone(input, ParallelInstruction.class); if (instruction.getRead() != null) { Source cloudSource = instruction.getRead().getSource(); cloudSource.setCodec(forCodec(cloudSource.getCodec(), replaceWithByteArrayCoder)); } else if (instruction.getWrite() != null) { com.google.api.services.dataflow.model.Sink cloudSink = instruction.getWrite().getSink(); cloudSink.setCodec(forCodec(cloudSink.getCodec(), replaceWithByteArrayCoder)); } else if (instruction.getParDo() != null) { instruction.setParDo( forParDoInstruction(instruction.getParDo(), replaceWithByteArrayCoder)); } else if (instruction.getPartialGroupByKey() != null) { PartialGroupByKeyInstruction pgbk = instruction.getPartialGroupByKey(); pgbk.setInputElementCodec(forCodec(pgbk.getInputElementCodec(), replaceWithByteArrayCoder)); } else if (instruction.getFlatten() != null) { } else { throw new RuntimeException("Unknown parallel instruction: " + input); } return instruction; } catch (IOException e) { throw new RuntimeException( String.format( "Failed to replace unknown coder with " + "LengthPrefixCoder for : {%s}", input), e); } }  static MutableNetwork<Node, Edge> andReplaceForRunnerNetwork(
      MutableNetwork<Node, Edge> network); static MutableNetwork<Node, Edge> forSdkNetwork(MutableNetwork<Node, Edge> network); static List<SideInputInfo> forSideInputInfos(
      List<SideInputInfo> sideInputInfos, boolean replaceWithByteArrayCoder);  }
@Test public void testLengthPrefixReadInstructionCoder() throws Exception { ReadInstruction readInstruction = new ReadInstruction(); readInstruction.setSource( new Source() .setCodec(CloudObjects.asCloudObject(windowedValueCoder, null))); instruction.setRead(readInstruction); ParallelInstruction prefixedInstruction = forParallelInstruction(instruction, false); assertEqualsAsJson( CloudObjects.asCloudObject(prefixedWindowedValueCoder, null), prefixedInstruction.getRead().getSource().getCodec()); assertEqualsAsJson( readInstruction.getSource().getCodec(), CloudObjects.asCloudObject(windowedValueCoder, null)); }
@Test public void testLengthPrefixWriteInstructionCoder() throws Exception { WriteInstruction writeInstruction = new WriteInstruction(); writeInstruction.setSink( new Sink() .setCodec(CloudObjects.asCloudObject(windowedValueCoder, null))); instruction.setWrite(writeInstruction); ParallelInstruction prefixedInstruction = forParallelInstruction(instruction, false); assertEqualsAsJson( CloudObjects.asCloudObject(prefixedWindowedValueCoder, null), prefixedInstruction.getWrite().getSink().getCodec()); assertEqualsAsJson( CloudObjects.asCloudObject(windowedValueCoder, null), writeInstruction.getSink().getCodec()); }
@Test public void testLengthPrefixParDoInstructionCoder() throws Exception { ParDoInstruction parDo = new ParDoInstruction(); CloudObject spec = CloudObject.forClassName(MERGE_BUCKETS_DO_FN); spec.put( WorkerPropertyNames.INPUT_CODER, CloudObjects.asCloudObject(windowedValueCoder, null)); parDo.setUserFn(spec); instruction.setParDo(parDo); ParallelInstruction prefixedInstruction = forParallelInstruction(instruction, false); assertEqualsAsJson( CloudObjects.asCloudObject(prefixedWindowedValueCoder, null), prefixedInstruction.getParDo().getUserFn().get(WorkerPropertyNames.INPUT_CODER)); assertEqualsAsJson( CloudObjects.asCloudObject(windowedValueCoder, null), parDo.getUserFn().get(WorkerPropertyNames.INPUT_CODER)); }