public static Environment createOrGetDefaultEnvironment(String type, String config) { if (Strings.isNullOrEmpty(type)) { return JAVA_SDK_HARNESS_ENVIRONMENT; } switch (type) { case ENVIRONMENT_EMBEDDED: return createEmbeddedEnvironment(config); case ENVIRONMENT_PROCESS: return createProcessEnvironment(config); case ENVIRONMENT_DOCKER: default: return createDockerEnvironment(config); } }
@Test public void testToFromProtoBounded() throws Exception { // TODO: Split into two tests. assumeThat(source, instanceOf(BoundedSource.class)); BoundedSource<?> boundedSource = (BoundedSource<?>) this.source; Read.Bounded<?> boundedRead = Read.from(boundedSource); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); ReadPayload payload = ReadTranslation.toProto(boundedRead, components); assertThat(payload.getIsBounded(), equalTo(RunnerApi.IsBounded.Enum.BOUNDED)); BoundedSource<?> deserializedSource = ReadTranslation.boundedSourceFromProto(payload); assertThat(deserializedSource, equalTo(source)); }
@Test public void testToFromProtoUnbounded() throws Exception { assumeThat(source, instanceOf(UnboundedSource.class)); UnboundedSource<?, ?> unboundedSource = (UnboundedSource<?, ?>) this.source; Read.Unbounded<?> unboundedRead = Read.from(unboundedSource); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); ReadPayload payload = ReadTranslation.toProto(unboundedRead, components); assertThat(payload.getIsBounded(), equalTo(RunnerApi.IsBounded.Enum.UNBOUNDED)); UnboundedSource<?, ?> deserializedSource = ReadTranslation.unboundedSourceFromProto(payload); assertThat(deserializedSource, equalTo(source)); }
@Test public void testEnvironment() { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); Environment env = Environments.createDockerEnvironment("java_test"); String id = sdkComponents.registerEnvironment(env); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); Environment rehydratedEnv = rehydratedComponents.getEnvironment(id); assertThat(rehydratedEnv, equalTo(env)); assertThat(rehydratedComponents.getEnvironment(id), theInstance(rehydratedEnv)); } }
@Test public void testToProtoAndBack() throws Exception { WindowingStrategy<?, ?> windowingStrategy = toProtoAndBackSpec.getWindowingStrategy(); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); WindowingStrategy<?, ?> toProtoAndBackWindowingStrategy = WindowingStrategyTranslation.fromProto( WindowingStrategyTranslation.toMessageProto(windowingStrategy, components)); assertThat( toProtoAndBackWindowingStrategy, equalTo((WindowingStrategy) windowingStrategy.fixDefaults())); }
@Test public void testStateSpecToFromProto() throws Exception { // Encode SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.StateSpec stateSpecProto = ParDoTranslation.translateStateSpec(stateSpec, sdkComponents); // Decode RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); StateSpec<?> deserializedStateSpec = ParDoTranslation.fromProto(stateSpecProto, rehydratedComponents); assertThat(stateSpec, equalTo(deserializedStateSpec)); } }
@Test public void testToProto() throws Exception { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); ParDoPayload payload = ParDoTranslation.translateParDo(parDo, p, components); assertThat(ParDoTranslation.getDoFn(payload), equalTo(parDo.getFn())); assertThat(ParDoTranslation.getMainOutputTag(payload), equalTo(parDo.getMainOutputTag())); for (PCollectionView<?> view : parDo.getSideInputs()) { payload.getSideInputsOrThrow(view.getTagInternal().getId()); } }
@Test public void testEncodedProto() throws Exception { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.TestStreamPayload payload = TestStreamTranslation.payloadForTestStream(testStream, components); verifyTestStreamEncoding( testStream, payload, RehydratedComponents.forComponents(components.toComponents())); }
@Test public void testSimpleCoder() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); Coder<?> coder = VarIntCoder.of(); String id = sdkComponents.registerCoder(coder); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); Coder<?> rehydratedCoder = rehydratedComponents.getCoder(id); assertThat(rehydratedCoder, equalTo((Coder) coder)); assertThat(rehydratedComponents.getCoder(id), theInstance((Coder) rehydratedCoder)); }
@Test public void testViewFnTranslation() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); assertEquals( new TestViewFn(), PCollectionViewTranslation.viewFnFromProto( ParDoTranslation.translateViewFn(new TestViewFn(), sdkComponents))); }
@Test public void testToFromProto() throws InvalidProtocolBufferException { pipeline.apply(GenerateSequence.from(0)).apply(Window.<Long>into((WindowFn) windowFn)); final AtomicReference<AppliedPTransform<?, ?, Assign<?>>> assign = new AtomicReference<>(null); pipeline.traverseTopologically( new PipelineVisitor.Defaults() { @Override public void visitPrimitiveTransform(Node node) { if (node.getTransform() instanceof Window.Assign) { checkState(assign.get() == null); assign.set( (AppliedPTransform<?, ?, Assign<?>>) node.toAppliedPTransform(getPipeline())); } } }); checkState(assign.get() != null); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); WindowIntoPayload payload = WindowIntoTranslation.toProto(assign.get().getTransform(), components); assertEquals(windowFn, WindowingStrategyTranslation.windowFnFromProto(payload.getWindowFn())); }
@Test public void testEncodedProto() throws Exception { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.WriteFilesPayload payload = WriteFilesTranslation.payloadForWriteFiles(writeFiles, components); assertThat( payload.getRunnerDeterminedSharding(), equalTo( writeFiles.getNumShardsProvider() == null && writeFiles.getComputeNumShards() == null)); assertThat(payload.getWindowedWrites(), equalTo(writeFiles.getWindowedWrites())); assertThat( (FileBasedSink<String, Void, String>) WriteFilesTranslation.sinkFromProto(payload.getSink()), equalTo(writeFiles.getSink())); }
@Test public void testEncodeDecodeFields() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.PCollection protoCollection = PCollectionTranslation.toProto(testCollection, sdkComponents); RehydratedComponents protoComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); Coder<?> decodedCoder = protoComponents.getCoder(protoCollection.getCoderId()); WindowingStrategy<?, ?> decodedStrategy = protoComponents.getWindowingStrategy(protoCollection.getWindowingStrategyId()); IsBounded decodedIsBounded = PCollectionTranslation.isBounded(protoCollection); assertThat(decodedCoder, equalTo(testCollection.getCoder())); assertThat(decodedStrategy, equalTo(testCollection.getWindowingStrategy().fixDefaults())); assertThat(decodedIsBounded, equalTo(testCollection.isBounded())); }
@Test public void getEnvironmentUnknownFnType() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN) .build()) .build(); Optional<Environment> env = Environments.getEnvironment(builder, rehydratedComponents); assertThat(env.isPresent(), is(false)); }
@Test public void testEncodeDecodeCycle() throws Exception { // Encode SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.PCollection protoCollection = PCollectionTranslation.toProto(testCollection, sdkComponents); RehydratedComponents protoComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); // Decode Pipeline pipeline = Pipeline.create(); PCollection<?> decodedCollection = PCollectionTranslation.fromProto(protoCollection, pipeline, protoComponents); // Verify assertThat(decodedCollection.getCoder(), equalTo(testCollection.getCoder())); assertThat( decodedCollection.getWindowingStrategy(), equalTo(testCollection.getWindowingStrategy().fixDefaults())); assertThat(decodedCollection.isBounded(), equalTo(testCollection.isBounded())); }
@Test public void testWindowMappingFnTranslation() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); assertEquals( new GlobalWindows().getDefaultWindowMappingFn(), PCollectionViewTranslation.windowMappingFnFromProto( ParDoTranslation.translateWindowMappingFn( new GlobalWindows().getDefaultWindowMappingFn(), sdkComponents))); }
@Test public void testToProtoAndBackWithComponents() throws Exception { WindowingStrategy<?, ?> windowingStrategy = toProtoAndBackSpec.getWindowingStrategy(); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.WindowingStrategy proto = WindowingStrategyTranslation.toProto(windowingStrategy, components); RehydratedComponents protoComponents = RehydratedComponents.forComponents(components.toComponents()); assertThat( WindowingStrategyTranslation.fromProto(proto, protoComponents).fixDefaults(), equalTo(windowingStrategy.fixDefaults())); protoComponents.getCoder( components.registerCoder(windowingStrategy.getWindowFn().windowCoder())); assertThat( proto.getAssignsToOneWindow(), equalTo(windowingStrategy.getWindowFn().assignsToOneWindow())); } }
@Test public void testRegistrarEncodedProto() throws Exception { PCollection<String> output = p.apply(testStream); AppliedPTransform<PBegin, PCollection<String>, TestStream<String>> appliedTestStream = AppliedPTransform.of("fakeName", PBegin.in(p).expand(), output.expand(), testStream, p); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.FunctionSpec spec = PTransformTranslation.toProto(appliedTestStream, components).getSpec(); assertThat(spec.getUrn(), equalTo(TEST_STREAM_TRANSFORM_URN)); RunnerApi.TestStreamPayload payload = TestStreamPayload.parseFrom(spec.getPayload()); verifyTestStreamEncoding( testStream, payload, RehydratedComponents.forComponents(components.toComponents())); }
@Test public void testWindowingStrategy() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); WindowingStrategy windowingStrategy = WindowingStrategy.of(FixedWindows.of(Duration.millis(1))) .withAllowedLateness(Duration.standardSeconds(4)); String id = sdkComponents.registerWindowingStrategy(windowingStrategy); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); WindowingStrategy<?, ?> rehydratedStrategy = rehydratedComponents.getWindowingStrategy(id); assertThat(rehydratedStrategy, equalTo((WindowingStrategy) windowingStrategy.fixDefaults())); assertThat( rehydratedComponents.getWindowingStrategy(id), theInstance((WindowingStrategy) rehydratedStrategy)); }
@Before public void setup() { partialComponents = Components.newBuilder() .putTransforms( "impulse", PTransform.newBuilder() .setUniqueName("Impulse") .putOutputs("output", "impulse.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) .build()) .putPcollections("impulse.out", pc("impulse.out")) .putEnvironments("go", Environments.createDockerEnvironment("go")) .putEnvironments("py", Environments.createDockerEnvironment("py")) .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) .putWindowingStrategies( "ws", WindowingStrategy.newBuilder().setWindowCoderId("windowCoder").build()) .build(); }