This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/texera.git
The following commit(s) were added to refs/heads/main by this push:
new 88ca47f2f4 test(amber): replace commented-out PythonWorkflowWorkerSpec
with pythonworker proxy unit tests (#7488)
88ca47f2f4 is described below
commit 88ca47f2f4d4ae26535012951b9c64cb9af3c91c
Author: Xinyuan Lin <[email protected]>
AuthorDate: Sun Aug 9 18:57:22 2026 -0700
test(amber): replace commented-out PythonWorkflowWorkerSpec with
pythonworker proxy unit tests (#7488)
### What changes were proposed in this PR?
#7447 proposed deleting `PythonWorkflowWorkerSpec.scala`, which has been
fully commented out for years and no longer compiles against today's
APIs. Following the review feedback there
(https://github.com/apache/texera/pull/7447#issuecomment-5234671357),
this PR replaces the dead file with real unit tests instead of only
deleting it.
The commented-out spec drove `PythonWorkflowWorker` end-to-end, which
needs a live Python process — that path is covered by the e2e tests.
What can be unit-tested without Python is the JVM side of the JVM↔Python
Arrow Flight bridge, which had no coverage until now:
```
JVM Python
┌───────────────────────────────────────┐
│ PythonWorkflowWorker (actor) │ e2e-tested only (needs
Python)
│ ├─ PythonProxyClient ──── Flight ────┼────▶ network_receiver.py
│ │ PythonProxyClientSpec: fake │
│ │ Python Flight server in Scala │
│ └─ PythonProxyServer ◀─── Flight ────┼───── network_sender.py
│ PythonProxyServerSpec: test │
│ plays the Python Flight client │
└───────────────────────────────────────┘
```
| New spec | Subject | Behavior pinned down |
|---|---|---|
| `PythonProxyServerSpec` | `PythonProxyServer` / `AmberProducer` |
`handshake` completes the port promise and replies `ok`; `control`
actions route `ControlInvocation` / `ReturnInvocation` to the output
gateway on the control channel and ack with a little-endian credit
value; `Data` / `State` / `ECM` puts are reassembled into `DataFrame` /
`StateFrame` (loop envelope preserved) / `EmbeddedControlMessage` and
acked with credits |
| `PythonProxyClientSpec` | `PythonProxyClient` | heartbeat handshake
happens before the queue is drained; queued `ControlInvocation` /
`ReturnInvocation` / actor commands arrive as `control` / `actor`
actions with intact protobuf payloads; `DataFrame` / `StateFrame` / ECM
puts arrive under the right `PythonDataHeader` with tuples, loop
envelope, and bytes intact; queue-size acks update `getQueuedCredit`;
connection retries abort with `WorkflowRuntimeException` (no server
listening, non-`ack` heartbeat); `close()` before any connection does
not throw |
Both specs stand in for the Python worker with plain Arrow Flight
components (`FlightClient` / `NoOpFlightProducer`), so no Python process
is involved. Together with the existing `WorkerBatchInternalQueueSpec`
and `PythonWorkflowWorkerStartupConfigSpec`, every class in the
`pythonworker` package except the actor itself now has unit coverage.
The commented-out `PythonWorkflowWorkerSpec.scala` is removed,
superseded by these tests.
### Any related issues, documentation, discussions?
Supersedes #7447.
### How was this PR tested?
This PR is tests-only. Both new specs pass locally:
```
sbt "WorkflowExecutionService/testOnly
org.apache.texera.amber.engine.architecture.pythonworker.PythonProxyServerSpec
org.apache.texera.amber.engine.architecture.pythonworker.PythonProxyClientSpec"
```
### Was this PR authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Fable 5)
---
.../pythonworker/PythonProxyClientSpec.scala | 379 +++++++++++++++++++++
.../pythonworker/PythonProxyServerSpec.scala | 306 +++++++++++++++++
.../pythonworker/PythonWorkflowWorkerSpec.scala | 201 -----------
3 files changed, 685 insertions(+), 201 deletions(-)
diff --git
a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClientSpec.scala
b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClientSpec.scala
new file mode 100644
index 0000000000..f536b61d1c
--- /dev/null
+++
b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClientSpec.scala
@@ -0,0 +1,379 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.pythonworker
+
+import com.twitter.util.Promise
+import org.apache.arrow.flight.{
+ Action,
+ FlightProducer,
+ FlightServer,
+ FlightStream,
+ Location,
+ NoOpFlightProducer,
+ PutResult,
+ Result
+}
+import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.vector.VarBinaryVector
+import org.apache.texera.amber.core.WorkflowRuntimeException
+import org.apache.texera.amber.core.state.State
+import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema,
Tuple}
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import
org.apache.texera.amber.engine.architecture.pythonworker.WorkerBatchInternalQueue.{
+ DataElement,
+ EmbeddedControlMessageElement
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ ControlInvocation,
+ EmbeddedControlMessage,
+ EmbeddedControlMessageType,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{
+ EmptyReturn,
+ ReturnInvocation
+}
+import org.apache.texera.amber.engine.common.actormessage.{CreditUpdate,
PythonActorMessage}
+import org.apache.texera.amber.engine.common.ambermessage.{
+ DataFrame,
+ DirectControlMessagePayloadV2,
+ PythonControlMessage,
+ PythonDataHeader,
+ StateFrame
+}
+import org.apache.texera.amber.util.ArrowUtils
+import org.scalatest.flatspec.AnyFlatSpec
+
+import java.net.ServerSocket
+import java.nio.charset.StandardCharsets
+import java.util.concurrent.ConcurrentLinkedQueue
+import scala.jdk.CollectionConverters._
+
+/**
+ * Exercises PythonProxyClient without a Python process: a Scala
+ * NoOpFlightProducer stands in for the Python worker's Flight server
+ * (`network_receiver.py`), speaking the same action/put protocol — heartbeat
+ * acks, per-action queue-size replies, and per-put credit metadata.
+ */
+class PythonProxyClientSpec extends AnyFlatSpec {
+
+ private val clientWorkerId = ActorVirtualIdentity("python-proxy-client")
+ private val controllerId = ActorVirtualIdentity("CONTROLLER")
+ private val upstreamChannel =
+ ChannelIdentity(controllerId, clientWorkerId, isControl = true)
+
+ private def freePort: Int = {
+ val socket = new ServerSocket(0)
+ try socket.getLocalPort
+ finally socket.close()
+ }
+
+ private def quietly(f: => Unit): Unit =
+ try f
+ catch { case _: Exception => () }
+
+ private def awaitTrue(timeoutMs: Long = 10000)(cond: => Boolean): Unit = {
+ val deadline = System.currentTimeMillis() + timeoutMs
+ while (!cond && System.currentTimeMillis() < deadline) Thread.sleep(20)
+ assert(cond, s"condition not met within ${timeoutMs}ms")
+ }
+
+ //
---------------------------------------------------------------------------
+ // Stand-in for the Python worker's Flight server: records every action and
+ // put it receives and acks them the way network_receiver.py does.
+ //
---------------------------------------------------------------------------
+
+ private case class RecordedPut(
+ header: PythonDataHeader,
+ tuples: Vector[Tuple],
+ ecmBytes: Option[Array[Byte]]
+ )
+
+ private class FakePythonWorkerServer(heartbeatBody: String) extends
NoOpFlightProducer {
+ val actionsReceived = new ConcurrentLinkedQueue[(String, Array[Byte])]()
+ val putsReceived = new ConcurrentLinkedQueue[RecordedPut]()
+ @volatile var reportedQueueSize: Long = 0L
+
+ override def doAction(
+ context: FlightProducer.CallContext,
+ action: Action,
+ listener: FlightProducer.StreamListener[Result]
+ ): Unit = {
+ actionsReceived.add((action.getType, action.getBody))
+ action.getType match {
+ case "heartbeat" =>
+ listener.onNext(new
Result(heartbeatBody.getBytes(StandardCharsets.UTF_8)))
+ listener.onCompleted()
+ case "control" | "actor" =>
+ // The ack body is the Python worker's queue size, sent as a decimal
+ // string; the client parses it for flow control.
+ listener.onNext(
+ new
Result(reportedQueueSize.toString.getBytes(StandardCharsets.UTF_8))
+ )
+ listener.onCompleted()
+ case _ => // "shutdown": the client does not consume a reply
+ listener.onCompleted()
+ }
+ }
+
+ override def acceptPut(
+ context: FlightProducer.CallContext,
+ flightStream: FlightStream,
+ ackStream: FlightProducer.StreamListener[PutResult]
+ ): Runnable = { () =>
+ {
+ val header =
PythonDataHeader.parseFrom(flightStream.getDescriptor.getCommand)
+ val tuples = Vector.newBuilder[Tuple]
+ var ecmBytes: Option[Array[Byte]] = None
+ while (flightStream.next()) {
+ val root = flightStream.getRoot
+ if (header.payloadType == "ECM") {
+ ecmBytes =
Some(root.getVector("payload").asInstanceOf[VarBinaryVector].get(0))
+ } else {
+ (0 until root.getRowCount).foreach(i => tuples +=
ArrowUtils.getTexeraTuple(i, root))
+ }
+ }
+ putsReceived.add(RecordedPut(header, tuples.result(), ecmBytes))
+ // Ack with the queue size as put metadata, the way the Python side
does.
+ val ackAllocator = new RootAllocator(1024)
+ try {
+ val buf = ackAllocator.buffer(java.lang.Long.BYTES)
+ buf.writeLong(reportedQueueSize)
+ ackStream.onNext(PutResult.metadata(buf))
+ buf.close()
+ } finally {
+ ackAllocator.close()
+ }
+ }
+ }
+ }
+
+ //
---------------------------------------------------------------------------
+ // Test harness — fake server + a PythonProxyClient main loop on its own
thread
+ //
---------------------------------------------------------------------------
+
+ private class ClientFixture(heartbeatBody: String = "ack") {
+ val producer = new FakePythonWorkerServer(heartbeatBody)
+ val serverAllocator = new RootAllocator()
+ val port: Int = freePort
+ val server: FlightServer = FlightServer
+ .builder(serverAllocator, Location.forGrpcInsecure("localhost", port),
producer)
+ .build()
+ server.start()
+ val portPromise: Promise[Int] = Promise[Int]()
+ portPromise.setValue(port)
+ val client = new PythonProxyClient(portPromise, clientWorkerId)
+ private val loopThread = new Thread(() => client.run(),
"python-proxy-client-test-loop")
+ loopThread.setDaemon(true)
+ // The tear-down interrupt below surfaces as an InterruptedException out of
+ // the blocking queue take; keep it off stderr.
+ loopThread.setUncaughtExceptionHandler((_, _) => ())
+
+ def start(): Unit = loopThread.start()
+
+ def close(): Unit = {
+ quietly(client.close())
+ loopThread.interrupt() // unblock the main loop's queue take
+ loopThread.join(5000)
+ quietly(server.close())
+ quietly(serverAllocator.close())
+ }
+ }
+
+ private def withClient(heartbeatBody: String = "ack")(test: ClientFixture =>
Unit): Unit = {
+ val fixture = new ClientFixture(heartbeatBody)
+ try test(fixture)
+ finally fixture.close()
+ }
+
+ //
---------------------------------------------------------------------------
+ // connection + control path
+ //
---------------------------------------------------------------------------
+
+ "PythonProxyClient" should
+ "connect via heartbeat and forward an enqueued ControlInvocation as a
control action" in {
+ withClient() { fixture =>
+ fixture.start()
+ val invocation = ControlInvocation(
+ "OpenExecutor",
+ EmptyRequest(),
+ AsyncRPCContext(controllerId, clientWorkerId),
+ 5L
+ )
+ fixture.client.enqueueCommand(invocation, upstreamChannel)
+
+ awaitTrue()(fixture.producer.actionsReceived.asScala.exists(_._1 ==
"control"))
+ assert(
+ fixture.producer.actionsReceived.asScala.head._1 == "heartbeat",
+ "the client must handshake before draining its queue"
+ )
+ val body = fixture.producer.actionsReceived.asScala.find(_._1 ==
"control").get._2
+ val parsed = PythonControlMessage.parseFrom(body)
+ assert(parsed.tag == upstreamChannel)
+ assert(
+ parsed.payload ==
+
DirectControlMessagePayloadV2.defaultInstance.withControlInvocation(invocation)
+ )
+ }
+ }
+
+ it should "forward an enqueued ReturnInvocation as a control action" in {
+ withClient() { fixture =>
+ fixture.start()
+ val reply = ReturnInvocation(42L, EmptyReturn())
+ fixture.client.enqueueCommand(reply, upstreamChannel)
+
+ awaitTrue()(fixture.producer.actionsReceived.asScala.exists(_._1 ==
"control"))
+ val body = fixture.producer.actionsReceived.asScala.find(_._1 ==
"control").get._2
+ val parsed = PythonControlMessage.parseFrom(body)
+ assert(parsed.tag == upstreamChannel)
+ assert(
+ parsed.payload ==
+
DirectControlMessagePayloadV2.defaultInstance.withReturnInvocation(reply)
+ )
+ }
+ }
+
+ it should "forward an enqueued actor command as an actor action" in {
+ withClient() { fixture =>
+ fixture.start()
+ fixture.client.enqueueActorCommand(CreditUpdate())
+
+ awaitTrue()(fixture.producer.actionsReceived.asScala.exists(_._1 ==
"actor"))
+ val body = fixture.producer.actionsReceived.asScala.find(_._1 ==
"actor").get._2
+ assert(PythonActorMessage.parseFrom(body) ==
PythonActorMessage(CreditUpdate()))
+ }
+ }
+
+ it should "update its Python-queue credit from the action ack" in {
+ withClient() { fixture =>
+ fixture.producer.reportedQueueSize = 7L
+ fixture.start()
+ assert(fixture.client.getQueuedCredit == 0L)
+
+ fixture.client.enqueueActorCommand(CreditUpdate())
+ awaitTrue()(fixture.client.getQueuedCredit == 7L)
+ }
+ }
+
+ //
---------------------------------------------------------------------------
+ // data path — Data / State / ECM puts
+ //
---------------------------------------------------------------------------
+
+ it should "stream an enqueued DataFrame to the server under a Data header"
in {
+ withClient() { fixture =>
+ fixture.start()
+ val schema = Schema()
+ .add(new Attribute("v", AttributeType.INTEGER))
+ .add(new Attribute("s", AttributeType.STRING))
+ def tuple(v: Int, s: String): Tuple =
+ Tuple.builder(schema).addSequentially(Array(Int.box(v), s)).build()
+ val tuples = Array(tuple(1, "a"), tuple(2, "b"))
+ val dataChannel = ChannelIdentity(clientWorkerId, controllerId,
isControl = false)
+
+ fixture.client.enqueueData(DataElement(DataFrame(tuples), dataChannel))
+
+ awaitTrue()(!fixture.producer.putsReceived.isEmpty)
+ val put = fixture.producer.putsReceived.peek()
+ assert(put.header == PythonDataHeader(dataChannel, "Data"))
+ assert(put.tuples == tuples.toVector)
+ }
+ }
+
+ it should "stream an enqueued StateFrame as a single State row carrying the
loop envelope" in {
+ withClient() { fixture =>
+ fixture.start()
+ val state = State(Map("count" -> "5"))
+ val dataChannel = ChannelIdentity(clientWorkerId, controllerId,
isControl = false)
+
+ fixture.client.enqueueData(
+ DataElement(StateFrame(state, loopCounter = 3L, loopStartId =
"loop-start-1"), dataChannel)
+ )
+
+ awaitTrue()(!fixture.producer.putsReceived.isEmpty)
+ val put = fixture.producer.putsReceived.peek()
+ assert(put.header == PythonDataHeader(dataChannel, "State"))
+ assert(put.tuples.size == 1)
+ val row = put.tuples.head
+ assert(State.fromTuple(row) == state)
+ assert(State.loopCounterFrom(row) == 3L)
+ assert(State.loopStartIdFrom(row) == "loop-start-1")
+ }
+ }
+
+ it should "stream an enqueued EmbeddedControlMessage as a serialized ECM
put" in {
+ withClient() { fixture =>
+ fixture.start()
+ val ecm = EmbeddedControlMessage(
+ EmbeddedControlMessageIdentity("ecm-1"),
+ EmbeddedControlMessageType.NO_ALIGNMENT,
+ Seq.empty,
+ Map.empty
+ )
+ val dataChannel = ChannelIdentity(clientWorkerId, controllerId,
isControl = false)
+
+ fixture.client.enqueueData(EmbeddedControlMessageElement(ecm,
dataChannel))
+
+ awaitTrue()(!fixture.producer.putsReceived.isEmpty)
+ val put = fixture.producer.putsReceived.peek()
+ assert(put.header == PythonDataHeader(dataChannel, "ECM"))
+ assert(put.ecmBytes.isDefined)
+ assert(EmbeddedControlMessage.parseFrom(put.ecmBytes.get) == ecm)
+ }
+ }
+
+ //
---------------------------------------------------------------------------
+ // failure paths
+ //
---------------------------------------------------------------------------
+
+ it should "abort with WorkflowRuntimeException once connection retries are
exhausted" in {
+ // A free port with nothing listening: every connection attempt is refused.
+ val portPromise = Promise[Int]()
+ portPromise.setValue(freePort)
+ val client = new PythonProxyClient(portPromise,
ActorVirtualIdentity("no-server"))
+ assertThrows[WorkflowRuntimeException] {
+ client.run()
+ }
+ }
+
+ it should "abort when the server heartbeat does not reply ack" in {
+ withClient(heartbeatBody = "nak") { fixture =>
+ // run() on the test thread: it must give up after the retry budget
+ // because the heartbeat body check fails on every attempt.
+ assertThrows[WorkflowRuntimeException] {
+ fixture.client.run()
+ }
+ }
+ }
+
+ it should "tolerate close() before any connection is established" in {
+ val portPromise = Promise[Int]()
+ portPromise.setValue(freePort)
+ val client = new PythonProxyClient(portPromise,
ActorVirtualIdentity("never-connected"))
+ // Must not throw: the internal null flight client is handled.
+ client.close()
+ succeed
+ }
+}
diff --git
a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyServerSpec.scala
b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyServerSpec.scala
new file mode 100644
index 0000000000..1e4c531a0b
--- /dev/null
+++
b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyServerSpec.scala
@@ -0,0 +1,306 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.pythonworker
+
+import com.twitter.util.{Await, Duration, Promise}
+import org.apache.arrow.flight.{Action, FlightClient, FlightDescriptor,
Location, SyncPutListener}
+import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema
=> ArrowSchema}
+import org.apache.arrow.vector.{VarBinaryVector, VectorSchemaRoot}
+import org.apache.texera.amber.core.state.State
+import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema,
Tuple}
+import org.apache.texera.amber.core.virtualidentity.{
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ EmbeddedControlMessageIdentity
+}
+import
org.apache.texera.amber.engine.architecture.messaginglayer.NetworkOutputGateway
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+ AsyncRPCContext,
+ ControlInvocation,
+ EmbeddedControlMessage,
+ EmbeddedControlMessageType,
+ EmptyRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{
+ EmptyReturn,
+ ReturnInvocation
+}
+import org.apache.texera.amber.engine.common.ambermessage.{
+ DataFrame,
+ DirectControlMessagePayloadV2,
+ PythonControlMessage,
+ PythonDataHeader,
+ StateFrame,
+ WorkflowFIFOMessage
+}
+import org.apache.texera.amber.util.ArrowUtils
+import org.scalatest.flatspec.AnyFlatSpec
+
+import java.nio.charset.StandardCharsets
+import java.nio.{ByteBuffer, ByteOrder}
+import java.util.concurrent.{ConcurrentLinkedQueue, TimeUnit}
+import scala.jdk.CollectionConverters._
+
+/**
+ * Exercises the JVM-side Flight server of the Python worker bridge without a
+ * Python process: the test plays the Python worker's part with a plain Arrow
+ * FlightClient and asserts on what the server forwards to its
+ * NetworkOutputGateway. This is the JVM half of the protocol implemented by
+ * `network_sender.py` / `network_receiver.py`.
+ */
+class PythonProxyServerSpec extends AnyFlatSpec {
+
+ private val jvmWorkerId = ActorVirtualIdentity("jvm-proxy-server")
+ private val pythonWorkerId = ActorVirtualIdentity("python-worker-1")
+ private val controllerId = ActorVirtualIdentity("CONTROLLER")
+
+ //
---------------------------------------------------------------------------
+ // Test harness — a running PythonProxyServer plus a FlightClient connected
to
+ // it, with every message the server emits captured from the output gateway.
+ //
---------------------------------------------------------------------------
+
+ private class ServerFixture {
+ val sentMessages = new ConcurrentLinkedQueue[WorkflowFIFOMessage]()
+ val gateway = new NetworkOutputGateway(jvmWorkerId, msg => {
sentMessages.add(msg); () })
+ val portPromise: Promise[Int] = Promise[Int]()
+ val server = new PythonProxyServer(gateway, jvmWorkerId, portPromise)
+ server.run() // FlightServer.start() binds and returns; gRPC threads serve
+ val clientAllocator = new RootAllocator()
+ val flightClient: FlightClient = FlightClient
+ .builder(
+ clientAllocator,
+ Location.forGrpcInsecure("localhost", server.getPortNumber.get())
+ )
+ .build()
+
+ def messages: List[WorkflowFIFOMessage] = sentMessages.asScala.toList
+
+ def close(): Unit = {
+ quietly(flightClient.close())
+ quietly(server.close())
+ quietly(clientAllocator.close())
+ }
+ }
+
+ private def quietly(f: => Unit): Unit =
+ try f
+ catch { case _: Exception => () }
+
+ private def withServer(test: ServerFixture => Unit): Unit = {
+ val fixture = new ServerFixture
+ try test(fixture)
+ finally fixture.close()
+ }
+
+ private def awaitTrue(timeoutMs: Long = 10000)(cond: => Boolean): Unit = {
+ val deadline = System.currentTimeMillis() + timeoutMs
+ while (!cond && System.currentTimeMillis() < deadline) Thread.sleep(20)
+ assert(cond, s"condition not met within ${timeoutMs}ms")
+ }
+
+ /** Mirror of PythonProxyClient.writeArrowStream: put one batch of rows. */
+ private def putBatch(
+ fixture: ServerFixture,
+ header: PythonDataHeader,
+ arrowSchema: ArrowSchema,
+ fillRoot: VectorSchemaRoot => Unit
+ ): Long = {
+ val descriptor = FlightDescriptor.command(header.toByteArray)
+ val listener = new SyncPutListener
+ val root = VectorSchemaRoot.create(arrowSchema, fixture.clientAllocator)
+ try {
+ val writer = fixture.flightClient.startPut(descriptor, root, listener)
+ root.allocateNew()
+ fillRoot(root)
+ writer.putNext()
+ root.clear()
+ writer.completed()
+ val ack = listener.poll(5, TimeUnit.SECONDS)
+ assert(ack != null, "no credit ack received for the put within 5s")
+ val ackBuf = ack.getApplicationMetadata
+ val credit = ackBuf.getLong(0)
+ ackBuf.close()
+ listener.close()
+ credit
+ } finally {
+ root.close()
+ }
+ }
+
+ //
---------------------------------------------------------------------------
+ // Port + handshake
+ //
---------------------------------------------------------------------------
+
+ "PythonProxyServer" should "bind a free local port exposed via
getPortNumber" in {
+ withServer { fixture =>
+ assert(fixture.server.getPortNumber.get() > 0)
+ }
+ }
+
+ it should "complete the port-number promise and reply ok on a handshake
action" in {
+ withServer { fixture =>
+ val result = fixture.flightClient
+ .doAction(new Action("handshake",
"45678".getBytes(StandardCharsets.UTF_8)))
+ .next()
+ assert(new String(result.getBody, StandardCharsets.UTF_8) == "ok")
+ assert(Await.result(fixture.portPromise, Duration.fromSeconds(5)) ==
45678)
+ }
+ }
+
+ //
---------------------------------------------------------------------------
+ // control actions — routed to the gateway on the control channel, acked with
+ // a little-endian credit value
+ //
---------------------------------------------------------------------------
+
+ it should "route a control action carrying a ControlInvocation to the tag's
receiver" in {
+ withServer { fixture =>
+ val tag = ChannelIdentity(pythonWorkerId, controllerId, isControl = true)
+ val invocation = ControlInvocation(
+ "StartWorker",
+ EmptyRequest(),
+ AsyncRPCContext(pythonWorkerId, controllerId),
+ 3L
+ )
+ val controlMessage = PythonControlMessage(
+ tag,
+
DirectControlMessagePayloadV2.defaultInstance.withControlInvocation(invocation)
+ )
+
+ val results = fixture.flightClient.doAction(new Action("control",
controlMessage.toByteArray))
+ val credit =
ByteBuffer.wrap(results.next().getBody).order(ByteOrder.LITTLE_ENDIAN).getLong
+ assert(credit == 30L, "the action ack must carry the credit value
little-endian")
+ assert(!results.hasNext, "a control action must produce exactly one
result")
+
+ val out = fixture.messages
+ assert(out.size == 1)
+ // The gateway sends on ITS control channel to the tag's receiver.
+ assert(out.head.channelId == ChannelIdentity(jvmWorkerId, controllerId,
isControl = true))
+ assert(out.head.payload == invocation)
+ }
+ }
+
+ it should "route a control action carrying a ReturnInvocation to the tag's
receiver" in {
+ withServer { fixture =>
+ val tag = ChannelIdentity(pythonWorkerId, controllerId, isControl = true)
+ val reply = ReturnInvocation(42L, EmptyReturn())
+ val controlMessage = PythonControlMessage(
+ tag,
+
DirectControlMessagePayloadV2.defaultInstance.withReturnInvocation(reply)
+ )
+
+ val results = fixture.flightClient.doAction(new Action("control",
controlMessage.toByteArray))
+ results.next()
+
+ val out = fixture.messages
+ assert(out.size == 1)
+ assert(out.head.channelId == ChannelIdentity(jvmWorkerId, controllerId,
isControl = true))
+ assert(out.head.payload == reply)
+ }
+ }
+
+ //
---------------------------------------------------------------------------
+ // puts — data / state / ECM payload types, acked with a credit value
+ //
---------------------------------------------------------------------------
+
+ it should "deliver a Data put as a DataFrame on the header's channel and ack
with credits" in {
+ withServer { fixture =>
+ val to =
+ ChannelIdentity(pythonWorkerId, ActorVirtualIdentity("downstream"),
isControl = false)
+ val schema = Schema()
+ .add(new Attribute("v", AttributeType.INTEGER))
+ .add(new Attribute("s", AttributeType.STRING))
+ def tuple(v: Int, s: String): Tuple =
+ Tuple.builder(schema).addSequentially(Array(Int.box(v), s)).build()
+ val tuples = Array(tuple(1, "a"), tuple(2, "b"))
+
+ val credit = putBatch(
+ fixture,
+ PythonDataHeader(to, "Data"),
+ ArrowUtils.fromTexeraSchema(schema),
+ root => tuples.foreach(t => ArrowUtils.appendTexeraTuple(t, root))
+ )
+ assert(credit == 31L, "the put ack must carry the credit value")
+
+ awaitTrue()(fixture.messages.nonEmpty)
+ val out = fixture.messages
+ assert(out.size == 1)
+ assert(out.head.channelId == to)
+ assert(out.head.payload == DataFrame(tuples))
+ }
+ }
+
+ it should "reassemble a single-row State put into a StateFrame carrying the
loop envelope" in {
+ withServer { fixture =>
+ val to =
+ ChannelIdentity(pythonWorkerId, ActorVirtualIdentity("downstream"),
isControl = false)
+ val state = State(Map("count" -> "5"))
+ val row = state.toTuple(loopCounter = 3L, loopStartId = "loop-start-1")
+
+ putBatch(
+ fixture,
+ PythonDataHeader(to, "State"),
+ ArrowUtils.fromTexeraSchema(State.schema),
+ root => ArrowUtils.appendTexeraTuple(row, root)
+ )
+
+ awaitTrue()(fixture.messages.nonEmpty)
+ val out = fixture.messages
+ assert(out.size == 1)
+ assert(out.head.channelId == to)
+ assert(out.head.payload == StateFrame(state, 3L, "loop-start-1"))
+ }
+ }
+
+ it should "parse an ECM put back into the original EmbeddedControlMessage"
in {
+ withServer { fixture =>
+ val to =
+ ChannelIdentity(pythonWorkerId, ActorVirtualIdentity("downstream"),
isControl = false)
+ val ecm = EmbeddedControlMessage(
+ EmbeddedControlMessageIdentity("ecm-1"),
+ EmbeddedControlMessageType.NO_ALIGNMENT,
+ Seq.empty,
+ Map.empty
+ )
+ // Mirror of PythonProxyClient.sendECM: one VarBinary "payload" column,
+ // one row holding the serialized ECM.
+ val field = new Field("payload", FieldType.nullable(new
ArrowType.Binary), null)
+ val arrowSchema = new ArrowSchema(List(field).asJava)
+
+ putBatch(
+ fixture,
+ PythonDataHeader(to, "ECM"),
+ arrowSchema,
+ root => {
+ val vector = root.getVector("payload").asInstanceOf[VarBinaryVector]
+ vector.setSafe(0, ecm.toByteArray)
+ vector.setValueCount(1)
+ root.setRowCount(1)
+ }
+ )
+
+ awaitTrue()(fixture.messages.nonEmpty)
+ val out = fixture.messages
+ assert(out.size == 1)
+ assert(out.head.channelId == to)
+ assert(out.head.payload == ecm)
+ }
+ }
+}
diff --git
a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerSpec.scala
b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerSpec.scala
deleted file mode 100644
index 9741dc4e45..0000000000
---
a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerSpec.scala
+++ /dev/null
@@ -1,201 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-//package org.apache.texera.amber.engine.architecture.pythonworker
-//
-//import org.apache.pekko.actor.{ActorRef, ActorSystem, Props}
-//import org.apache.pekko.testkit.{ImplicitSender, TestActorRef, TestKit}
-//import org.apache.texera.amber.clustering.SingleNodeListener
-//import
org.apache.texera.amber.engine.architecture.common.WorkflowActor.{NetworkAck,
NetworkMessage}
-//import
org.apache.texera.amber.engine.architecture.pythonworker.promisehandlers.InitializeOperatorLogicHandler.InitializeOperatorLogic
-//import
org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.OneToOnePartitioning
-//import
org.apache.texera.amber.engine.architecture.worker.controlcommands.LinkOrdinal
-//import
org.apache.texera.amber.engine.architecture.worker.promisehandlers.AddPartitioningHandler.AddPartitioning
-//import
org.apache.texera.amber.engine.architecture.worker.promisehandlers.OpenOperatorHandler.OpenOperator
-//import
org.apache.texera.amber.engine.architecture.worker.promisehandlers.UpdateInputLinkingHandler.UpdateInputLinking
-//import org.apache.texera.amber.engine.common.Constants
-//import org.apache.texera.amber.engine.common.ambermessage.{
-// ChannelIdentity,
-// ControlPayload,
-// DataFrame,
-// DataPayload,
-// EndOfUpstream,
-// WorkflowFIFOMessage
-//}
-//import
org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.{ControlInvocation,
ReturnInvocation}
-//import org.apache.texera.amber.engine.common.virtualidentity.util.COORDINATOR
-//import org.apache.texera.amber.core.virtualidentity.{
-// ActorVirtualIdentity,
-// PhysicalLink,
-// PhysicalLink,
-// OperatorIdentity
-//}
-//import org.apache.texera.amber.engine.e2e.TestOperators
-//import org.apache.texera.workflow.common.tuple.Tuple
-//import org.apache.texera.workflow.common.tuple.schema.{Attribute,
AttributeType, Schema}
-//import org.scalamock.scalatest.MockFactory
-//import org.scalatest.BeforeAndAfterAll
-//import org.scalatest.flatspec.AnyFlatSpecLike
-//
-//import scala.concurrent.duration.DurationInt
-//
-//class PythonWorkflowWorkerSpec
-// extends TestKit(ActorSystem("PythonWorkerSpec"))
-// with ImplicitSender
-// with AnyFlatSpecLike
-// with BeforeAndAfterAll
-// with MockFactory {
-//
-// override def beforeAll: Unit = {
-// system.actorOf(Props[SingleNodeListener], "cluster-info")
-// }
-// override def afterAll: Unit = {
-// TestKit.shutdownActorSystem(system)
-// }
-// private val identifier1 = ActorVirtualIdentity("worker-1")
-// private val identifier2 = ActorVirtualIdentity("worker-2")
-// private val operatorIdentity = OperatorIdentity("testWorkflow",
"testOperator")
-// private val layerId1 =
-// PhysicalLink(operatorIdentity.workflow, operatorIdentity.operator,
"1st-layer")
-// private val layerId2 =
-// PhysicalLink(operatorIdentity.workflow, operatorIdentity.operator,
"2nd-layer")
-// private val pythonOp = TestOperators.pythonOpDesc()
-// private val link = PhysicalLink(layerId1, 0, layerId2, 0)
-// private val schema = Schema
-// .newBuilder()
-// .add(new Attribute("text", AttributeType.STRING))
-// .build()
-// private val initialization = InitializeOperatorLogic(
-// pythonOp.code,
-// isSource = false,
-// Seq(LinkOrdinal(link, 0)),
-// Seq(LinkOrdinal(link, 0)),
-// schema
-// )
-//
-// def sendControlToWorker(
-// worker: ActorRef,
-// controls: Array[ControlInvocation],
-// beginSeqNum: Long = 0
-// ): Unit = {
-// var seq = beginSeqNum
-// controls.foreach { ctrl =>
-// worker ! NetworkMessage(
-// seq,
-// WorkflowFIFOMessage(ChannelIdentity(COORDINATOR, identifier1, true),
seq, ctrl)
-// )
-// val received = receiveWhile(3.seconds) {
-// case NetworkAck(id, credits) =>
-// // pass
-// case NetworkMessage(id, fifoPayload) =>
-// fifoPayload.payload.asInstanceOf[ControlPayload] match {
-// case ControlInvocation(commandID, command) => assert(commandID
== seq)
-// case ReturnInvocation(originalCommandID, controlReturn) =>
-// assert(originalCommandID == seq)
-// case _ => ???
-// }
-// worker ! NetworkAck(id,
Constants.unprocessedBatchesSizeLimitInBytesPerWorkerPair)
-// }
-// seq += 1
-// }
-// }
-//
-// def mkWorker: ActorRef = TestActorRef(new
PythonWorkflowWorker(identifier1))
-//
-// "python worker" should "start" in {
-// val worker = mkWorker
-// sendControlToWorker(worker, Array(ControlInvocation(0, initialization)))
-// }
-//
-// "python worker" should "process data" in {
-// val worker = mkWorker
-// sendControlToWorker(worker, Array(ControlInvocation(0, initialization)))
-// val mockPolicy = OneToOnePartitioning(1, Array(identifier2))
-// val openControl = ControlInvocation(1, OpenOperator())
-// val invocation = ControlInvocation(2, AddPartitioning(link, mockPolicy))
-// val updateInputLinking = ControlInvocation(3,
UpdateInputLinking(identifier2, link))
-// sendControlToWorker(worker, Array(openControl, invocation,
updateInputLinking), 1)
-// worker ! NetworkMessage(
-// 4,
-// WorkflowFIFOMessage(
-// ChannelIdentity(identifier2, identifier1, false),
-// 0,
-// DataFrame(
-// Array(
-// Tuple
-// .newBuilder(schema)
-// .add("text", AttributeType.STRING, "123")
-// .build()
-// )
-// )
-// )
-// )
-// expectMsgClass(classOf[NetworkAck])
-// val data = receiveOne(30.seconds)
-//
assert(data.asInstanceOf[NetworkMessage].internalMessage.payload.isInstanceOf[DataFrame])
-// }
-//
-// "python worker" should "process data and receive end marker" in {
-// val worker = mkWorker
-// sendControlToWorker(worker, Array(ControlInvocation(0, initialization)))
-// val mockPolicy = OneToOnePartitioning(100, Array(identifier2))
-// val openControl = ControlInvocation(1, OpenOperator())
-// val invocation = ControlInvocation(2, AddPartitioning(link, mockPolicy))
-// val updateInputLinking = ControlInvocation(3,
UpdateInputLinking(identifier2, link))
-// sendControlToWorker(worker, Array(openControl, invocation,
updateInputLinking), 1)
-// worker ! NetworkMessage(
-// 4,
-// WorkflowFIFOMessage(
-// ChannelIdentity(identifier2, identifier1, false),
-// 0,
-// DataFrame(
-// (0 until 100)
-// .map(_ =>
-// Tuple
-// .newBuilder(schema)
-// .add("text", AttributeType.STRING, "123")
-// .build()
-// )
-// .toArray
-// )
-// )
-// )
-// expectMsgClass(classOf[NetworkAck])
-// val data = receiveOne(30.seconds)
-//
assert(data.asInstanceOf[NetworkMessage].internalMessage.payload.isInstanceOf[DataFrame])
-// worker ! NetworkMessage(
-// 5,
-// WorkflowFIFOMessage(
-// ChannelIdentity(identifier2, identifier1, false),
-// 1,
-// EndOfUpstream()
-// )
-// )
-// expectMsgClass(classOf[NetworkAck])
-// receiveWhile(10.seconds) {
-// case NetworkMessage(id, fifoPayload) =>
-// fifoPayload.payload match {
-// case payload: ControlPayload => //skip
-// case payload: DataPayload =>
assert(payload.isInstanceOf[EndOfUpstream])
-// case _ => ???
-// }
-// }
-// }
-//
-//}
