(texera) 01/01: test(amber): cover Pekko actor references (#7219)

[gi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-7219-18ebb0d1e6dded6f8bc04b7c52be4fa31cb27916
in repository https://gitbox.apache.org/repos/asf/texera.git

commit 58179a08a73203bae5c3cd6e615b592a235713ea
Author: Xinyuan Lin <xi...@uci.edu>
AuthorDate: Sat Aug 1 00:01:44 2026 -0700

    test(amber): cover Pekko actor references (#7219)
    
    ### What changes were proposed in this PR?
    
    Adds focused unit coverage for Pekko actor-reference mapping: known and
    unknown credit forwarding, parent lookup de-duplication and reset, FIFO
    stashed-message delivery, coordinator registration notifications, and
    reverse lookup removal.
    
    ### Any related issues, documentation, discussions?
    
    Closes #7216
    
    ### How was this PR tested?
    
    - `WorkflowExecutionService/testOnly
    org.apache.texera.amber.engine.architecture.common.PekkoActorRefMappingServiceSpec`
    — 4 passed.
    - `WorkflowExecutionService/Test/scalafmtCheck`
    - `WorkflowExecutionService/Test/scalafix --check`
    
    Mutation proof (each production mutation was reverted):
    
    | Production mutation | Focused test result |
    | --- | --- |
    | Removed `stash.enqueue(msg)` from `forwardToActor` | 3 passed, 1
    failed |
    | Removed `queriedActorVirtualIdentities.add(id)` from lookup
    propagation | 3 passed, 1 failed |
    
    ### Was this PR authored or co-authored using generative AI tooling?
    
    Generated-by: Codex (GPT-5)
---
 .../common/PekkoActorRefMappingServiceSpec.scala   | 173 +++++++++++++++++++++
 1 file changed, 173 insertions(+)

diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/common/PekkoActorRefMappingServiceSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/common/PekkoActorRefMappingServiceSpec.scala
new file mode 100644
index 0000000000..6b5cf140ac
--- /dev/null
+++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/common/PekkoActorRefMappingServiceSpec.scala
@@ -0,0 +1,173 @@
+/*
+ * 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.common
+
+import org.apache.pekko.actor.{Actor, ActorSystem, Props}
+import org.apache.pekko.testkit.{TestActorRef, TestKit, TestProbe}
+import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
+import org.apache.texera.amber.engine.architecture.common.WorkflowActor.{
+  CreditRequest,
+  GetActorRef,
+  NetworkMessage,
+  RegisterActorRef
+}
+import org.apache.texera.amber.engine.common.ambermessage.{DataFrame, WorkflowFIFOMessage}
+import org.apache.texera.amber.engine.common.virtualidentity.util.COORDINATOR
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.flatspec.AnyFlatSpecLike
+
+import scala.concurrent.duration.DurationInt
+
+/**
+  * Unit tests for [[PekkoActorRefMappingService]].
+  *
+  * Like [[PekkoMessageTransferServiceSpec]], these tests use a minimal actor
+  * spawned by Pekko TestKit to supply the live [[ActorContext]] that
+  * [[PekkoActorService]] eagerly dereferences. Spawning it as a child of a
+  * [[TestProbe]] also makes the service's parent lookup observable.
+  */
+class PekkoActorRefMappingServiceSpec
+    extends TestKit(ActorSystem("PekkoActorRefMappingServiceSpec"))
+    with AnyFlatSpecLike
+    with BeforeAndAfterAll {
+
+  override def afterAll(): Unit = {
+    TestKit.shutdownActorSystem(system)
+  }
+
+  private val upstreamId = ActorVirtualIdentity("upstream")
+  private val workerId = ActorVirtualIdentity("mapping-service-worker")
+  private val contextCounter = new java.util.concurrent.atomic.AtomicInteger(0)
+
+  private def channelTo(destination: ActorVirtualIdentity): ChannelIdentity =
+    ChannelIdentity(
+      fromWorkerId = upstreamId,
+      toWorkerId = destination,
+      isControl = true
+    )
+
+  private def networkMessageTo(
+      destination: ActorVirtualIdentity,
+      messageId: Long,
+      sequenceNumber: Long
+  ): NetworkMessage =
+    NetworkMessage(
+      messageId,
+      WorkflowFIFOMessage(channelTo(destination), sequenceNumber, DataFrame(Array.empty))
+    )
+
+  private def newActorService(
+      id: ActorVirtualIdentity,
+      parent: TestProbe
+  ): PekkoActorService = {
+    val contextHolder =
+      TestActorRef[ActorRefMappingServiceContextHolder](
+        Props(new ActorRefMappingServiceContextHolder),
+        parent.ref,
+        s"actor-ref-mapping-context-${contextCounter.incrementAndGet()}"
+      )
+    new PekkoActorService(id, contextHolder.underlyingActor.context)
+  }
+
+  "askForCredit" should "forward a request only when the destination ref is known" in {
+    val parent = TestProbe()
+    val destination = TestProbe()
+    val unknownDestination = ActorVirtualIdentity("unknown-credit-destination")
+    val knownDestination = ActorVirtualIdentity("known-credit-destination")
+    val service = new PekkoActorRefMappingService(newActorService(workerId, parent))
+
+    service.askForCredit(channelTo(unknownDestination))
+    parent.expectNoMessage(100.millis)
+
+    service.registerActorRef(knownDestination, destination.ref)
+    val channel = channelTo(knownDestination)
+    service.askForCredit(channel)
+
+    destination.expectMsg(CreditRequest(channel))
+  }
+
+  "forwardToActor" should "stash unknown messages and ask the parent for their ref once" in {
+    val parent = TestProbe()
+    val destination = ActorVirtualIdentity("unknown-message-destination")
+    val service = new PekkoActorRefMappingService(newActorService(workerId, parent))
+    val first = networkMessageTo(destination, messageId = 1L, sequenceNumber = 0L)
+    val second = networkMessageTo(destination, messageId = 2L, sequenceNumber = 1L)
+
+    service.forwardToActor(first)
+    val lookup = parent.expectMsgType[GetActorRef]
+    assert(lookup.id == destination)
+    assert(lookup.replyTo == Set(service.self))
+
+    service.forwardToActor(second)
+    parent.expectNoMessage(100.millis)
+
+    service.clearQueriedActorRefs()
+    service.forwardToActor(networkMessageTo(destination, messageId = 3L, sequenceNumber = 2L))
+    assert(parent.expectMsgType[GetActorRef].id == destination)
+  }
+
+  "registerActorRef" should "drain stashed messages in FIFO order and notify coordinator waiters" in {
+    val parent = TestProbe()
+    val destination = ActorVirtualIdentity("coordinator-wait-destination")
+    val waiterOne = TestProbe()
+    val waiterTwo = TestProbe()
+    val registered = TestProbe()
+    val service = new PekkoActorRefMappingService(newActorService(COORDINATOR, parent))
+    val first = networkMessageTo(destination, messageId = 10L, sequenceNumber = 0L)
+    val second = networkMessageTo(destination, messageId = 11L, sequenceNumber = 1L)
+
+    service.forwardToActor(first)
+    service.forwardToActor(second)
+    service.retrieveActorRef(destination, Set(waiterOne.ref, waiterTwo.ref))
+
+    service.registerActorRef(destination, registered.ref)
+
+    registered.expectMsg(first)
+    registered.expectMsg(second)
+    waiterOne.expectMsg(RegisterActorRef(destination, registered.ref))
+    waiterTwo.expectMsg(RegisterActorRef(destination, registered.ref))
+  }
+
+  "retrieveActorRef and removeActorRef" should "reply for known refs and remove their reverse lookup" in {
+    val parent = TestProbe()
+    val destination = ActorVirtualIdentity("registered-destination")
+    val registered = TestProbe()
+    val waiter = TestProbe()
+    val service = new PekkoActorRefMappingService(newActorService(workerId, parent))
+
+    service.registerActorRef(destination, registered.ref)
+    service.retrieveActorRef(destination, Set(waiter.ref))
+
+    waiter.expectMsg(RegisterActorRef(destination, registered.ref))
+    assert(service.hasActorRef(destination))
+    assert(service.getActorRef(destination) == registered.ref)
+    assert(service.findActorVirtualIdentity(registered.ref).contains(destination))
+
+    service.removeActorRef(destination)
+
+    assert(!service.hasActorRef(destination))
+    assert(service.findActorVirtualIdentity(registered.ref).isEmpty)
+  }
+}
+
+/** Minimal actor used only to obtain a live [[ActorContext]] from Pekko TestKit. */
+class ActorRefMappingServiceContextHolder extends Actor {
+  override def receive: Receive = { case _ => () }
+}