Fix loopback (#34678)

* added the tests

* updated the test

* fixed the comments

Author: liferoad

sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Environments.java

@@ -219,20 +219,17 @@ public static Environment createDockerEnvironment(String dockerImageUrl) {
   }
 
   private static Environment createExternalEnvironment(String externalServiceAddress) {
-    if (externalServiceAddress.isEmpty()) {
-      throw new IllegalArgumentException(
-          String.format(
-              "External service address must not be empty (set it using '--environmentOptions=%s=...'?).",
-              externalServiceAddressOption));
+    // Create the payload builder. If the address is empty, the payload will be empty,
+    // acting as a placeholder for late binding. For example, in the LOOPBACK case,
+    // the address is populated by PortableRunner#run before this method is called.
+    ExternalPayload.Builder payloadBuilder = ExternalPayload.newBuilder();
+    if (!externalServiceAddress.isEmpty()) {
+      payloadBuilder.setEndpoint(
+          ApiServiceDescriptor.newBuilder().setUrl(externalServiceAddress).build());
     }
     return Environment.newBuilder()
         .setUrn(BeamUrns.getUrn(StandardEnvironments.Environments.EXTERNAL))
-        .setPayload(
-            ExternalPayload.newBuilder()
-                .setEndpoint(
-                    ApiServiceDescriptor.newBuilder().setUrl(externalServiceAddress).build())
-                .build()
-                .toByteString())
+        .setPayload(payloadBuilder.build().toByteString())
         .build();
   }
 

sdks/java/extensions/python/build.gradle

@@ -33,6 +33,7 @@ dependencies {
     testImplementation library.java.junit
     testImplementation library.java.hamcrest
     testImplementation project(":sdks:java:core").sourceSets.test.output
+    testImplementation project(path: ":runners:direct-java", configuration: "shadowTest")
     testRuntimeOnly library.java.slf4j_simple
 }
 

sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/PythonExternalTransformTest.java

@@ -27,6 +27,8 @@
 import org.apache.beam.model.pipeline.v1.ExternalTransforms;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.PortablePipelineOptions;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.schemas.SchemaTranslation;
 import org.apache.beam.sdk.schemas.logicaltypes.MicrosInstant;
@@ -371,4 +373,22 @@ public void generateKwargsFromMap() {
     assertEquals(12L, (long) receivedRow.getInt64("longField"));
     assertEquals(15.6, (double) receivedRow.getDouble("doubleField"), 0);
   }
+
+  @Test
+  public void testLoopbackEnvironmentWithPythonExternalTransform() {
+    PortablePipelineOptions options =
+        PipelineOptionsFactory.create().as(PortablePipelineOptions.class);
+    options.setDefaultEnvironmentType("LOOPBACK");
+
+    Pipeline p = Pipeline.create(options);
+
+    PCollection<String> output =
+        p.apply(Create.of(KV.of("A", "x"), KV.of("A", "y"), KV.of("B", "z")))
+            .apply(
+                PythonExternalTransform
+                    .<PCollection<KV<String, String>>, PCollection<KV<String, Iterable<String>>>>
+                        from("apache_beam.GroupByKey"))
+            .apply(Keys.create());
+    PAssert.that(output).containsInAnyOrder("A", "B");
+  }
 }