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-7157-07834709c7cc35370afbc6523b8b3624ad13290f in repository https://gitbox.apache.org/repos/asf/texera.git
commit 2388b8573836f3f1a84aa808ed7b3a999e3626bf Author: Xinyuan Lin <[email protected]> AuthorDate: Fri Jul 31 16:37:42 2026 -0700 test(amber): cover websocket session state (#7157) ### What changes were proposed in this PR? Adds focused unit coverage for websocket event delivery, workflow-service subscription replacement and cleanup, session-map removal, and computing-unit access state. ### Any related issues, documentation, discussions? Closes #7153 ### How was this PR tested? - `WorkflowExecutionService/testOnly org.apache.texera.web.SessionStateSpec` — 3 succeeded. - `WorkflowExecutionService/Test/scalafmtCheck` and `WorkflowExecutionService/Test/scalafix --check`. Mutation proof: | Production mutation | Result | | --- | --- | | Replaced the serialized websocket payload with an empty string | 1 test failed | | Stopped retaining the subscribed workflow service | 1 test failed | | Stopped removing a session from the session-state map | 1 test failed | Each mutation was reverted before the final green test run. ### Was this PR authored or co-authored using generative AI tooling? Generated-by: Codex (GPT-5) --------- Signed-off-by: Xinyuan Lin <[email protected]> Co-authored-by: Copilot Autofix powered by AI <[email protected]> --- .../org/apache/texera/web/SessionStateSpec.scala | 160 +++++++++++++++++++++ 1 file changed, 160 insertions(+) diff --git a/amber/src/test/scala/org/apache/texera/web/SessionStateSpec.scala b/amber/src/test/scala/org/apache/texera/web/SessionStateSpec.scala new file mode 100644 index 0000000000..34fe6250f5 --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/web/SessionStateSpec.scala @@ -0,0 +1,160 @@ +/* + * 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.web + +import io.reactivex.rxjava3.disposables.Disposable +import org.apache.texera.amber.core.virtualidentity.WorkflowIdentity +import org.apache.texera.amber.util.JSONUtils.objectMapper +import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum +import org.apache.texera.web.model.websocket.event.{TexeraWebSocketEvent, WorkflowStateEvent} +import org.apache.texera.web.service.WorkflowService +import org.scalamock.scalatest.MockFactory +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers + +import java.util.UUID +import java.util.concurrent.{Future => JFuture} +import javax.websocket.{RemoteEndpoint, Session} +import scala.collection.mutable.ArrayBuffer + +class SessionStateSpec extends AnyFlatSpec with Matchers with MockFactory { + + private class TrackingDisposable extends Disposable { + var disposeCalls = 0 + + override def dispose(): Unit = disposeCalls += 1 + + override def isDisposed: Boolean = disposeCalls > 0 + } + + private class TestWorkflowService(id: Long) extends WorkflowService(WorkflowIdentity(id), 1, 10) { + val workflowSubscription = new TrackingDisposable + val executionSubscription = new TrackingDisposable + var workflowConnectCalls = 0 + var executionConnectCalls = 0 + var disconnectCalls = 0 + + override def connect(onNext: TexeraWebSocketEvent => Unit): Disposable = { + workflowConnectCalls += 1 + workflowSubscription + } + + override def connectToExecution(onNext: TexeraWebSocketEvent => Unit): Disposable = { + executionConnectCalls += 1 + executionSubscription + } + + override def disconnect(): Unit = disconnectCalls += 1 + } + + private def messageCollectingSession(): (Session, ArrayBuffer[String]) = { + val messages = ArrayBuffer[String]() + val async = mock[RemoteEndpoint.Async] + (async + .sendText(_: String)) + .expects(*) + .onCall { (message: String) => + messages += message + null.asInstanceOf[JFuture[Void]] + } + .anyNumberOfTimes() + + val session = mock[Session] + (() => session.getAsyncRemote).expects().returning(async).anyNumberOfTimes() + (session, messages) + } + + private def removeStateIfPresent(sessionId: String): Unit = { + try { + SessionState.removeState(sessionId) + } catch { + case _: NoSuchElementException => + } + } + + "SessionState" should "send websocket events as typed JSON messages" in { + val (session, messages) = messageCollectingSession() + val state = new SessionState(session) + + state.send(WorkflowStateEvent("RUNNING")) + + messages should have size 1 + val payload = objectMapper.readTree(messages.head) + payload.get("type").asText() shouldBe "WorkflowStateEvent" + payload.get("state").asText() shouldBe "RUNNING" + } + + it should "replace subscriptions before attaching a new workflow service" in { + val state = new SessionState(stub[Session]) + val firstService = new TestWorkflowService(1L) + val secondService = new TestWorkflowService(2L) + + state.subscribe(firstService) + state.getCurrentWorkflowState shouldBe Some(firstService) + firstService.workflowConnectCalls shouldBe 1 + firstService.executionConnectCalls shouldBe 1 + + state.subscribe(secondService) + firstService.workflowSubscription.disposeCalls shouldBe 1 + firstService.executionSubscription.disposeCalls shouldBe 1 + firstService.disconnectCalls shouldBe 1 + state.getCurrentWorkflowState shouldBe Some(secondService) + secondService.workflowConnectCalls shouldBe 1 + secondService.executionConnectCalls shouldBe 1 + + state.unsubscribe() + secondService.workflowSubscription.disposeCalls shouldBe 1 + secondService.executionSubscription.disposeCalls shouldBe 1 + secondService.disconnectCalls shouldBe 1 + state.getCurrentWorkflowState shouldBe None + } + + it should "remove registered sessions, cleaning up subscriptions, and retain their computing-unit access level" in { + val sessionId = UUID.randomUUID().toString + val state = new SessionState(stub[Session]) + val service = new TestWorkflowService(3L) + SessionState.setState(sessionId, state) + try { + SessionState.getState(sessionId) shouldBe state + SessionState.getAllSessionStates should contain(state) + + state.getUserComputingUnitAccess shouldBe PrivilegeEnum.NONE + state.setUserComputingUnitAccess(PrivilegeEnum.WRITE) + state.getUserComputingUnitAccess shouldBe PrivilegeEnum.WRITE + + state.subscribe(service) + service.workflowConnectCalls shouldBe 1 + service.executionConnectCalls shouldBe 1 + + SessionState.removeState(sessionId) + service.workflowSubscription.disposeCalls shouldBe 1 + service.executionSubscription.disposeCalls shouldBe 1 + service.disconnectCalls shouldBe 1 + + // state object should retain its access level even after being removed from the registry + state.getUserComputingUnitAccess shouldBe PrivilegeEnum.WRITE + + SessionState.getAllSessionStates should not contain state + a[NoSuchElementException] should be thrownBy SessionState.getState(sessionId) + } finally { + removeStateIfPresent(sessionId) + } + } +}
