@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)); }
public static SdkComponents create(PipelineOptions options) { SdkComponents sdkComponents = new SdkComponents(); PortablePipelineOptions portablePipelineOptions = options.as(PortablePipelineOptions.class); sdkComponents.registerEnvironment( Environments.createOrGetDefaultEnvironment( portablePipelineOptions.getDefaultEnvironmentType(), portablePipelineOptions.getDefaultEnvironmentConfig())); return sdkComponents; }
@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)); }
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); } }
public Optional<Environment> getEnvironment(PTransformNode parDo) { return Environments.getEnvironment(parDo.getId(), components); }
private static Environment createProcessEnvironment(String config) { try { ProcessPayloadReferenceJSON payloadReferenceJSON = MAPPER.readValue(config, ProcessPayloadReferenceJSON.class); return createProcessEnvironment( payloadReferenceJSON.getOs(), payloadReferenceJSON.getArch(), payloadReferenceJSON.getCommand(), payloadReferenceJSON.getEnv()); } catch (IOException e) { throw new RuntimeException( String.format("Unable to parse process environment config: %s", config), e); } }
@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 getEnvironmentRead() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); ReadPayload payload = ReadTranslation.toProto(Read.from(CountingSource.unbounded()), components); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN) .setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env, equalTo( components .toComponents() .getEnvironmentsOrThrow(payload.getSource().getEnvironmentId()))); }
@Test public void createEnvironments() throws IOException { assertThat( Environments.createOrGetDefaultEnvironment(Environments.ENVIRONMENT_DOCKER, "java"), is( Environment.newBuilder() .build())); assertThat( Environments.createOrGetDefaultEnvironment( Environments.ENVIRONMENT_PROCESS, "{\"os\": \"linux\", \"arch\": \"amd64\", \"command\": \"run.sh\", \"env\":{\"k1\": \"v1\", \"k2\": \"v2\"} }"), .build())); assertThat( Environments.createOrGetDefaultEnvironment( Environments.ENVIRONMENT_PROCESS, "{\"command\": \"run.sh\"}"), is(
@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 getEnvironmentCombine() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); CombinePayload payload = CombinePayload.newBuilder() .setCombineFn(CombineTranslation.toProto(Sum.ofLongs(), components)) .build(); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN) .setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env, equalTo( components .toComponents() .getEnvironmentsOrThrow(payload.getCombineFn().getEnvironmentId()))); } }
@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 getEnvironmentParDo() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); ParDoPayload payload = ParDoTranslation.translateParDo( .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env,
@Test public void toAndFromProto() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.Coder coderProto = CoderTranslation.toProto(coder, sdkComponents); Components encodedComponents = sdkComponents.toComponents(); Coder<?> decodedCoder = CoderTranslation.fromProto( coderProto, RehydratedComponents.forComponents(encodedComponents)); assertThat(decodedCoder, equalTo(coder)); if (KNOWN_CODERS.contains(coder)) { for (RunnerApi.Coder encodedCoder : encodedComponents.getCodersMap().values()) { assertThat( encodedCoder.getSpec().getSpec().getUrn(), not(equalTo(CoderTranslation.JAVA_SERIALIZED_CODER_URN))); } } }
@Test public void getEnvironmentWindowIntoCustom() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); WindowIntoPayload payload = WindowIntoPayload.newBuilder() .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env,
@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 getEnvironmentWindowIntoKnown() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); WindowIntoPayload payload = WindowIntoPayload.newBuilder() .setWindowFn( WindowingStrategyTranslation.toProto( FixedWindows.of(Duration.standardMinutes(5L)), components)) .build(); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env, equalTo( components .toComponents() .getEnvironmentsOrThrow(payload.getWindowFn().getEnvironmentId()))); }
@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 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 testViewFnTranslation() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); assertEquals( new TestViewFn(), PCollectionViewTranslation.viewFnFromProto( ParDoTranslation.translateViewFn(new TestViewFn(), sdkComponents))); }