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-7676-4b526eb26afcce34467d50b6ff597f4f06642dde in repository https://gitbox.apache.org/repos/asf/texera.git
commit ae04417af8a7da46c5b3c2a3d8da67c7065b58f4 Author: Xinyuan Lin <[email protected]> AuthorDate: Sat Aug 15 03:34:36 2026 +0000 test(amber): cover the websocket handshake and session wiring (#7676) ### What changes were proposed in this PR? `ServletAwareConfigurator` had **no spec**, and it is the handshake hook both websocket endpoints declare via `@ServerEndpoint(configurator = ...)` -- it decides who the user is and what computing-unit privilege they carry for every connection. `WorkflowWebsocketResource` consumes exactly what it writes, and its `myOnOpen` was entirely uncovered. Covering them together is what makes the pair meaningful: the tests hand the configurator's real output to the resource rather than a hand-built map. Measured with only these two specs running, so the numbers are attributable to them alone: | File | Lines before | Lines after | Branches after | |---|---|---|---| | `ServletAwareConfigurator.scala` | 0/27 (0%) | **27/27 (100%)** | 16/18 | | `WorkflowWebsocketResource.scala` | 53.2% | **61/62 (98.4%)** | 26/34 | Tests **6 -> 22**. `ServletAwareConfiguratorSpec` is new (7 tests); `WorkflowWebsocketResourceSpec` gains 9. Covered: both handshake modes and the four-way header guard leg by leg, the JWT query-parameter path with a token minted in-process, both failure arms and the exact partial state each leaves behind, `myOnOpen`'s workflow/computing-unit/privilege binding and its deliberate event ordering, and `myOnMsg`'s pagination, modify-logic, unrecognised-command and failure-mapping arms. ### Verification 32 mutations applied one at a time and reverted, production diff empty after each. **All 32 red, no survivors.** 18 during the build; 14 more added by review, every one of which was first *run* to confirm it survived before any test was changed, then re-run to prove the fix goes red on the intended test by name (read from the JUnit XML, since the console log does not name tests). The review round is where the real work was. Examples of what it caught: | Weakness | Why it passed | Fix | |---|---|---| | the pagination arm asserted only the response type | the stub returned a canned value, so a rewritten `pageIndex` shipped green | assert the stub received the same requestID/operatorID/pageIndex it was sent | | `numWorkers` asserted by value | `0` is the literal initializer of `numWorkerNodesInCluster` | assert event presence and order | | the repeated-header case | `headOption` vs `lastOption` was indistinguishable | feed a genuinely repeated header | ### On amber's test parallelism Three findings hinged on whether sibling amber suites run concurrently, so review measured it instead of arguing: three probe suites each sleeping 4s while bumping a shared high-water mark ran in 12.3s with a concurrent maximum of **1**. `amber/build.sbt:52` (`concurrentRestrictions in Global += Tags.limit(Tags.Test, 1)`) does serialize suites within this project. The probe file was deleted afterwards. Suites still share one JVM, so the spec restores every global it touches. ### Deliberately not included One partial branch in `myOnOpen` needs a real cluster. Several defects found along the way are reported rather than pinned, so a fix is not blocked by a test asserting the current behaviour: - `WorkflowWebsocketResource.scala:89` computes `executionStateOpt` **before** the try, so a failure raised while the execution is being created goes to the socket instead of the new execution's metadata store, where the frontend's error panel reads it. *This one is pinned*, as characterization, with a comment saying so. - `:98-106` and `:124-126` both test `workflowStateOpt` where they mean `executionStateOpt`, so a workflow with no execution NPEs instead of reporting "not initialized". - `:81` puts `objectMapper.readValue` outside the try, so an unparseable frame escapes the error mapper and the client is told nothing. - `ServletAwareConfigurator.scala:62`'s `.getOrElse("")` yields a privilege string `PrivilegeEnum.valueOf` cannot parse; unreachable today, but `getOrElse(PrivilegeEnum.NONE.name())` would be safer. - `WorkflowService.getOrCreate` keys its cache on the workflow id alone, so a second opener's `computingUnitId` is silently dropped. The cached arm is covered by asserting *identity*; the dropped id is described in a comment but **not** asserted, so fixing it will not break this spec. No production file is touched. ### Any related issues, documentation, discussions? Closes #7675 ### How was this PR tested? ``` STORAGE_ICEBERG_CATALOG_TYPE=postgres sbt "WorkflowExecutionService/testOnly org.apache.texera.web.ServletAwareConfiguratorSpec org.apache.texera.web.resource.WorkflowWebsocketResourceSpec" ``` ``` [info] Total number of tests run: 22 [info] Tests: succeeded 22, failed 0, canceled 0, ignored 0, pending 0 [info] All tests passed. ``` `Test/scalafmtCheck` and `Test/scalafix --check` both pass. ### Was this PR authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5) --- .../texera/web/ServletAwareConfiguratorSpec.scala | 254 +++++++++++ .../resource/WorkflowWebsocketResourceSpec.scala | 462 ++++++++++++++++++++- 2 files changed, 703 insertions(+), 13 deletions(-) diff --git a/amber/src/test/scala/org/apache/texera/web/ServletAwareConfiguratorSpec.scala b/amber/src/test/scala/org/apache/texera/web/ServletAwareConfiguratorSpec.scala new file mode 100644 index 0000000000..7623dd992e --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/web/ServletAwareConfiguratorSpec.scala @@ -0,0 +1,254 @@ +/* + * 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 ch.qos.logback.classic.{Level, Logger => LogbackLogger} +import org.apache.texera.auth.JwtAuth +import org.apache.texera.auth.util.HeaderField +import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum +import org.apache.texera.dao.jooq.generated.tables.pojos.User +import org.jose4j.jwt.JwtClaims +import org.scalamock.scalatest.MockFactory +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers + +import javax.websocket.server.{HandshakeRequest, ServerEndpointConfig} +import scala.jdk.CollectionConverters._ + +/** + * Unit tests for the websocket handshake configurator, which is the single place a websocket + * connection acquires its identity. It runs in two mutually exclusive modes: + * + * - KUBERNETES, selected when all four `x-user-*` headers are present. Envoy is trusted to have + * authenticated the caller, so the user is built straight from header values. + * - SINGLE NODE otherwise, where the identity comes from an `access-token` JWT in the query + * string and the computing-unit privilege is hard-coded to WRITE. + * + * Both modes write into `ServerEndpointConfig.getUserProperties`, which is the map + * `WorkflowWebsocketResource.myOnOpen`/`myOnMsg` later read. Every assertion below therefore reads + * values back OUT of that map rather than trusting an intermediate. + * + * Tokens are minted with the production `JwtAuth.jwtToken`, so they verify against the very + * consumer the configurator uses, whatever `auth.jwt.256-bit-secret` resolves to in this JVM (the + * same trick `JwtParserSpec` uses). No socket, container, database or filesystem is involved. + * + * Deliberately not asserted: the exact wording of the info/debug log lines. They are not a + * contract, and pinning them would break on any rewording. (Worth reporting rather than testing: + * the info line writes the user's e-mail address into the log.) + */ +class ServletAwareConfiguratorSpec extends AnyFlatSpec with Matchers with MockFactory { + + private val userId = "77" + private val userName = "k8s-subject" + private val userEmail = "[email protected]" + + /** All four headers the kubernetes branch requires, each value distinct from the others. */ + private val k8sHeaders: Map[String, Seq[String]] = Map( + // READ, never WRITE: WRITE is what the single-node branch hard-codes, so a test that fed WRITE + // here could not tell the two branches apart. + HeaderField.UserComputingUnitAccess -> Seq(PrivilegeEnum.READ.name()), + HeaderField.UserId -> Seq(userId), + HeaderField.UserName -> Seq(userName), + HeaderField.UserEmail -> Seq(userEmail) + ) + + private def handshake( + headers: Map[String, Seq[String]] = Map.empty, + queryString: String = "" + ): (ServerEndpointConfig, HandshakeRequest, java.util.Map[String, AnyRef]) = { + val javaHeaders = headers.map { case (name, values) => name -> values.asJava }.asJava + val request = mock[HandshakeRequest] + (() => request.getHeaders).expects().returning(javaHeaders).anyNumberOfTimes() + (() => request.getQueryString).expects().returning(queryString).anyNumberOfTimes() + + val properties = new java.util.HashMap[String, AnyRef]() + val config = mock[ServerEndpointConfig] + (() => config.getUserProperties).expects().returning(properties).anyNumberOfTimes() + + (config, request, properties) + } + + private def userIn(properties: java.util.Map[String, AnyRef]): User = + properties.get(classOf[User].getName).asInstanceOf[User] + + private def tokenFor(subject: String, uid: Int, email: String): String = { + val claims = new JwtClaims + claims.setSubject(subject) + claims.setClaim("userId", uid) + claims.setClaim("email", email) + // The shared consumer is built with setRequireExpirationTime(); an exp-less token is rejected. + claims.setExpirationTimeMinutesInTheFuture(10f) + JwtAuth.jwtToken(claims) + } + + /** + * amber ships no logback-test.xml, so `amber/src/main/resources/logback.xml` governs this JVM + * and pins `org.apache` at WARN — which leaves the info/debug lines of the kubernetes branch + * unexecuted. Raising this one logger for the duration of a single call executes them. The + * previous level is restored in `finally` because every amber suite shares one JVM. + */ + private def withDebugLogging[T](body: => T): T = { + val logger = org.slf4j.LoggerFactory + .getLogger(classOf[ServletAwareConfigurator]) + .asInstanceOf[LogbackLogger] + val previousLevel = logger.getLevel + logger.setLevel(Level.DEBUG) + try body + finally logger.setLevel(previousLevel) + } + + // -- kubernetes mode ---------------------------------------------------------- + + "modifyHandshake" should "build the user from the trusted headers in kubernetes mode" in { + val (config, request, properties) = handshake(headers = k8sHeaders) + + withDebugLogging { + new ServletAwareConfigurator().modifyHandshake(config, request, null) + } + + // Read back by key: the privilege and the user live under two different keys, and both + // `myOnOpen` (privilege) and `myOnMsg` (user) look them up by exactly these names. + properties.get(HeaderField.UserComputingUnitAccess) shouldBe PrivilegeEnum.READ.name() + val user = userIn(properties) + // Three mutually distinct values, so a setName/setEmail transposition cannot pass. + user.getUid.intValue() shouldBe userId.toInt + user.getName shouldBe userName + user.getEmail shouldBe userEmail + } + + it should "take the first value of a repeated header, not the last" in { + // HTTP headers are legitimately multi-valued, and `_.asScala.headOption` is what decides WHICH + // occurrence wins. In kubernetes mode that choice decides the identity and the privilege a + // websocket is granted, so a duplicate appended after Envoy's own must not be able to override + // it. `k8sHeaders` gives every header exactly one value, where head and last agree, so nothing + // else here can see the difference. + val (config, request, properties) = handshake(headers = + k8sHeaders ++ Map( + HeaderField.UserId -> Seq(userId, "99"), + HeaderField.UserComputingUnitAccess -> Seq( + PrivilegeEnum.READ.name(), + PrivilegeEnum.WRITE.name() + ) + ) + ) + + new ServletAwareConfigurator().modifyHandshake(config, request, null) + + properties.get(HeaderField.UserComputingUnitAccess) shouldBe PrivilegeEnum.READ.name() + userIn(properties).getUid.intValue() shouldBe userId.toInt + } + + it should "write an empty privilege string that myOnOpen cannot parse when the access header carries no value" in { + // CHARACTERIZATION OF A DEFECT, NOT A CONTRACT — and a strict fix SHOULD turn this test red. + // + // Reachability: the enclosing guard only checks that the key is PRESENT, so an empty value list + // still selects the kubernetes branch and the privilege falls through to `getOrElse("")`. No + // servlet container produces that shape: a header that is present always carries at least one + // value, `""` at worst, which yields `Some("")` rather than None. The fallback is therefore + // defensive dead code, reachable only by handing `getHeaders` an empty list directly as below. + // + // What it produces is not a graceful default but a deferred crash: `PrivilegeEnum` is a Java + // enum, so `WorkflowWebsocketResource.myOnOpen`'s `PrivilegeEnum.valueOf` on this value throws + // straight out of the @OnOpen handler. The last assertion records that consequence, so that + // changing the fallback to `PrivilegeEnum.NONE.name()` reads as the fix it would be. + val (config, request, properties) = + handshake(headers = k8sHeaders + (HeaderField.UserComputingUnitAccess -> Seq.empty)) + + new ServletAwareConfigurator().modifyHandshake(config, request, null) + + properties.get(HeaderField.UserComputingUnitAccess) shouldBe "" + // The user is still built: the fallback must not abort the branch. + userIn(properties).getName shouldBe userName + an[IllegalArgumentException] should be thrownBy PrivilegeEnum.valueOf( + properties.get(HeaderField.UserComputingUnitAccess).asInstanceOf[String] + ) + } + + it should "fall back to single-node mode when any one of the four headers is missing" in { + // One case per header, because a single "none of them present" case only proves the + // conjunction is false — it would stay green with any one of the four `contains` calls deleted. + k8sHeaders.keys.foreach { missing => + withClue(s"with $missing missing: ") { + val (config, request, properties) = + handshake(headers = k8sHeaders - missing, queryString = "wid=1&cuid=2") + + new ServletAwareConfigurator().modifyHandshake(config, request, null) + + // WRITE is the single-node constant; the surviving headers say READ, so the value in the + // map names which branch ran. (With `x-user-id` or `x-user-name`/`x-user-email` missing the + // kubernetes branch would instead throw before writing anything, leaving the map empty.) + properties.get(HeaderField.UserComputingUnitAccess) shouldBe PrivilegeEnum.WRITE.name() + properties.containsKey(classOf[User].getName) shouldBe false + } + } + } + + // -- single-node mode --------------------------------------------------------- + + it should "build the user from the access-token query parameter in single-node mode" in { + // The token is surrounded by other parameters on purpose: with it alone in the query string a + // naive "the query string IS the token" implementation would pass too. + val token = tokenFor("token-subject", 4242, "[email protected]") + val (config, request, properties) = + handshake(queryString = s"wid=1&access-token=$token&cuid=2") + + new ServletAwareConfigurator().modifyHandshake(config, request, null) + + properties.get(HeaderField.UserComputingUnitAccess) shouldBe PrivilegeEnum.WRITE.name() + val user = userIn(properties) + // The uid claim survives a JSON round trip as a Long, which is why the production cast is to + // Long and not Integer; asserting the Int value pins the narrowing too. + user.getUid.intValue() shouldBe 4242 + user.getName shouldBe "token-subject" + user.getEmail shouldBe "[email protected]" + } + + // -- the catch arm ------------------------------------------------------------ + + it should "swallow a malformed user-id header, leaving the properties untouched" in { + // "no exception escaped" on its own would pass with the whole method body deleted, so the + // assertion that carries the weight is the PARTIAL state the failure leaves behind: the parse + // of `x-user-id` happens before any write, so nothing at all reaches the map. + val (config, request, properties) = + handshake(headers = k8sHeaders + (HeaderField.UserId -> Seq("not-a-number"))) + + noException should be thrownBy + new ServletAwareConfigurator().modifyHandshake(config, request, null) + + properties.keySet.asScala shouldBe empty + } + + it should "swallow a tampered access token, keeping the privilege it already granted" in { + // The other half of the partial-state contract: the single-node branch writes the privilege + // BEFORE it parses the token, so a rejected token leaves that entry behind and no user. + val token = tokenFor("token-subject", 4242, "[email protected]") + val parts = token.split('.') + parts.length shouldBe 3 + val tampered = s"${parts(0)}.${parts(1)}.${parts(2).reverse}" + val (config, request, properties) = + handshake(queryString = s"wid=1&access-token=$tampered&cuid=2") + + noException should be thrownBy + new ServletAwareConfigurator().modifyHandshake(config, request, null) + + properties.get(HeaderField.UserComputingUnitAccess) shouldBe PrivilegeEnum.WRITE.name() + properties.containsKey(classOf[User].getName) shouldBe false + } +} diff --git a/amber/src/test/scala/org/apache/texera/web/resource/WorkflowWebsocketResourceSpec.scala b/amber/src/test/scala/org/apache/texera/web/resource/WorkflowWebsocketResourceSpec.scala index 20e0afadc1..b059c9ea43 100644 --- a/amber/src/test/scala/org/apache/texera/web/resource/WorkflowWebsocketResourceSpec.scala +++ b/amber/src/test/scala/org/apache/texera/web/resource/WorkflowWebsocketResourceSpec.scala @@ -19,19 +19,39 @@ package org.apache.texera.web.resource +import com.google.protobuf.timestamp.Timestamp +import org.apache.texera.amber.clustering.ClusterListener +import org.apache.texera.amber.core.tuple.{Attribute, AttributeType} import org.apache.texera.amber.core.virtualidentity.WorkflowIdentity +import org.apache.texera.amber.core.workflow.WorkflowContext +import org.apache.texera.amber.core.workflowruntimestate.FatalErrorType.{ + COMPILATION_ERROR, + EXECUTION_FAILURE +} +import org.apache.texera.amber.core.workflowruntimestate.WorkflowFatalError +import org.apache.texera.amber.operator.limit.LimitOpDesc import org.apache.texera.amber.util.JSONUtils.objectMapper import org.apache.texera.auth.util.HeaderField import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum import org.apache.texera.dao.jooq.generated.tables.pojos.User import org.apache.texera.web.SessionState -import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent +import org.apache.texera.web.model.websocket.event.{PaginatedResultEvent, TexeraWebSocketEvent} import org.apache.texera.web.model.websocket.request.{ HeartBeatRequest, + ModifyLogicRequest, + ResultPaginationRequest, + RetryRequest, TexeraWebSocketRequest, WorkflowExecuteRequest } -import org.apache.texera.web.service.WorkflowService +import org.apache.texera.web.model.websocket.response.ModifyLogicResponse +import org.apache.texera.web.service.{ + ExecutionReconfigurationService, + ExecutionResultService, + WorkflowExecutionService, + WorkflowService +} +import org.apache.texera.web.storage.{ExecutionStateStore, WorkflowStateStore} import org.scalamock.scalatest.MockFactory import org.scalatest.BeforeAndAfterEach import org.scalatest.flatspec.AnyFlatSpec @@ -40,11 +60,12 @@ import org.scalatest.matchers.should.Matchers import io.reactivex.rxjava3.disposables.Disposable import java.net.URI +import java.time.Instant import java.util.UUID import java.util.concurrent.{Future => JFuture} import javax.websocket.{RemoteEndpoint, Session} import scala.collection.mutable.ArrayBuffer -import scala.jdk.CollectionConverters.IteratorHasAsScala +import scala.jdk.CollectionConverters.{IteratorHasAsScala, MapHasAsJava, SeqHasAsJava} /** * Unit tests for the websocket endpoint's message handling. @@ -72,11 +93,13 @@ import scala.jdk.CollectionConverters.IteratorHasAsScala * - `myOnOpen`'s missing-`wid`/`cuid` and bogus-privilege paths, which fail with NPE / * IndexOutOfBounds / IllegalArgumentException. A harmless improvement (a default, or a clear * message) would break such a test. - * - `ModifyLogicRequest`, which reaches `executionReconfigurationService` — null until - * `executeWorkflow()` has run, so the only way to exercise it here is via an NPE that - * production never reaches. - * - `ResultPaginationRequest`, whose payload line needs a DB, and whose no-workflow case is a - * discarded `Option.foreach` — "nothing sent, nothing thrown" asserts nothing. + * - `ModifyLogicRequest` with a workflow but no execution: `executionService.getValue` is null, + * so the handler NPEs instead of reporting "workflow execution is not initialized". Same guard + * gap as `case other` below — reported, not pinned. It is also the only input that would tell + * `if (workflowStateOpt.isDefined)` apart from `if (executionStateOpt.isDefined)`, so that + * guard's exact condition is left unpinned on purpose; see the no-workflow case below. + * - `ResultPaginationRequest`'s real payload, which needs a DB; only the request/response + * pass-through is asserted here, against a stubbed result service. * - the deserialization line, already owned by `TexeraWebSocketRequestSpec`. */ class WorkflowWebsocketResourceSpec @@ -95,20 +118,61 @@ class WorkflowWebsocketResourceSpec // SessionState's registry is a JVM-global map. Remove only the ids this test registered, so the // suite cannot disturb (or be disturbed by) SessionStateSpec running in the same JVM. // - // Best-effort: `removeState` throws NoSuchElementException on an id that is already gone, and the + // Best-effort: `getState` throws NoSuchElementException on an id that is already gone, and the // myOnClose test removes its own entry by design. + // + // This runs the real production teardown — `removeState` disposes the session's rx subscriptions + // and calls `WorkflowService.disconnect` — rather than swapping in a workflow-less stand-in + // first. That matters: a swap removes the registry entry but leaves the ORIGINAL state's + // subscription undisposed, so the mock session stays reachable from the workflow's state store + // for the life of the JVM. The reason a swap looked necessary is that `disconnect` drops the user + // count to zero and reaches `AmberRuntime.scheduleCallThroughActorSystem`, whose actor system is + // null in a unit JVM; the myOnOpen tests take a user count of their own on the service they + // create precisely so that this call never reaches zero. See `holdWorkflowService`. override protected def afterEach(): Unit = { - registeredSessions.foreach(id => scala.util.Try(SessionState.removeState(id))) + registeredSessions.foreach { id => + scala.util.Try(SessionState.removeState(id)) + // Checked, not assumed. `removeState` unsubscribes BEFORE it removes, so anything that throws + // in `unsubscribe` — a user count reaching zero and touching the null actor system being the + // one that bites here — leaves the mock session in the JVM-global registry for the rest of + // the run, where a later suite's cluster broadcast would call an expired mock. + withClue(s"session $id was still registered after afterEach: ") { + scala.util.Try(SessionState.getState(id)).isFailure shouldBe true + } + } registeredSessions.clear() } + /** + * Returns the shared `WorkflowService` for `wid`, creating it if needed, and takes a user count + * on it that is never released. + * + * Two things follow, both wanted. `myOnOpen` must JOIN this instance rather than build its own, + * which is what makes the identity assertions below meaningful; and the count never reaches zero + * in `afterEach`, so the real teardown can run without touching a null actor system. + * + * The cost is one `WorkflowService` per wid left in `WorkflowService.workflowServiceMapping` with + * a user count of 1 for the life of the JVM. Nothing test-scoped is reachable from it once the + * session is removed, and the wid namespace 9701xx is reserved for this suite — note that + * `getOrCreate` keys that map on the workflow id ALONE (`mkWorkflowStateId`), so a sibling suite + * reusing one of these ids would silently get this instance back. + */ + private def holdWorkflowService(wid: Long, cuid: Int): WorkflowService = { + val service = WorkflowService.getOrCreate(WorkflowIdentity(wid), cuid) + service.lifeCycleManager.increaseUserCount() + service + } + /** * A mocked session. `sent` collects everything the endpoint writes back, which is the only * observable for most of these handlers. */ private def mockSession( id: String = UUID.randomUUID().toString, - uid: Option[Int] = None + uid: Option[Int] = None, + cuAccess: String = PrivilegeEnum.WRITE.name(), + wid: Long = 1L, + cuid: Int = 1 ): (Session, ArrayBuffer[String]) = { val sent = ArrayBuffer[String]() @@ -123,17 +187,22 @@ class WorkflowWebsocketResourceSpec .anyNumberOfTimes() val properties = new java.util.HashMap[String, Object]() - properties.put(HeaderField.UserComputingUnitAccess, PrivilegeEnum.WRITE.name()) + properties.put(HeaderField.UserComputingUnitAccess, cuAccess) uid.foreach { u => val user = new User() user.setUid(Integer.valueOf(u)) properties.put(classOf[User].getName, user) } + // What `myOnOpen` parses the workflow and computing-unit ids out of. + val parameters: java.util.Map[String, java.util.List[String]] = + Map("wid" -> Seq(wid.toString).asJava, "cuid" -> Seq(cuid.toString).asJava).asJava + val session = mock[Session] (() => session.getId).expects().returning(id).anyNumberOfTimes() (() => session.getAsyncRemote).expects().returning(async).anyNumberOfTimes() (() => session.getUserProperties).expects().returning(properties).anyNumberOfTimes() + (() => session.getRequestParameterMap).expects().returning(parameters).anyNumberOfTimes() (() => session.getRequestURI) .expects() .returning(new URI("ws://localhost/wsapi/workflow-websocket")) @@ -159,6 +228,16 @@ class WorkflowWebsocketResourceSpec private class TestWorkflowService(id: Long) extends WorkflowService(WorkflowIdentity(id), 1, 10) { var initCalls: List[(WorkflowExecuteRequest, Option[User], URI)] = Nil + /** + * When set, `initExecutionService` publishes this execution and then fails — the shape a real + * `executeWorkflow()` failure takes, and the only way to reach the endpoint's catch arm with an + * execution that did NOT exist when the frame arrived. + */ + var publishThenFail: Option[WorkflowExecutionService] = None + + val results = new StubResultService(workflowId, computingUnitId, stateStore) + override val resultService: ExecutionResultService = results + override def connect(onNext: TexeraWebSocketEvent => Unit): Disposable = Disposable.empty() override def connectToExecution(onNext: TexeraWebSocketEvent => Unit): Disposable = Disposable.empty() @@ -168,7 +247,97 @@ class WorkflowWebsocketResourceSpec req: WorkflowExecuteRequest, userOpt: Option[User], sessionUri: URI - ): Unit = initCalls = initCalls :+ ((req, userOpt, sessionUri)) + ): Unit = { + initCalls = initCalls :+ ((req, userOpt, sessionUri)) + publishThenFail.foreach { execution => + executionService.onNext(execution) + throw new IllegalStateException("failed while starting the execution") + } + } + } + + /** + * The real `handleResultPagination` resolves the latest execution out of the database and opens + * an Iceberg document; the endpoint's own contribution is only that the request travels in and + * the returned event travels back out. The echoed fields make both directions observable. + * + * The table and schema carry values on purpose. Production ships a companion + * `PaginatedResultEvent(req, table, schema)` that rebuilds requestID/operatorID/pageIndex FROM the + * request, so an endpoint that called the service and then constructed its own event would match + * on those three fields anyway. These two payload fields are the only ones that can only have + * come from the returned event. + */ + private class StubResultService( + workflowId: WorkflowIdentity, + computingUnitId: Int, + stateStore: WorkflowStateStore + ) extends ExecutionResultService(workflowId, computingUnitId, stateStore) { + var captured: Option[ResultPaginationRequest] = None + + override def handleResultPagination(request: ResultPaginationRequest): TexeraWebSocketEvent = { + captured = Some(request) + PaginatedResultEvent( + request.requestID, + request.operatorID, + request.pageIndex, + List(objectMapper.createObjectNode().put("c", 5)), + List(new Attribute("c", AttributeType.INTEGER)) + ) + } + } + + /** + * Production only builds an `ExecutionReconfigurationService` inside `executeWorkflow()`, but + * the field it lands in is a public var, so a test can put a stub there directly. The two + * `register*` overrides are the seams the service already ships for constructing it without a + * live `AmberClient` (see `ExecutionReconfigurationServiceSpec`). + */ + private class CapturingReconfigurationService(stateStore: ExecutionStateStore) + extends ExecutionReconfigurationService(client = null, stateStore, workflow = null) { + var captured: Option[ModifyLogicRequest] = None + + override protected def registerWorkerCompletionCallback(): Unit = () + override protected def registerCompletionDiffHandler(): Unit = () + + override def modifyOperatorLogic(request: ModifyLogicRequest): TexeraWebSocketEvent = { + captured = Some(request) + // Both fields are derived from the request that arrived, so the response the endpoint + // forwards proves which request the service actually saw. + ModifyLogicResponse( + request.operator.operatorIdentifier.id, + isValid = true, + errorMessage = request.operator.asInstanceOf[LimitOpDesc].limit.toString + ) + } + } + + /** + * A real execution, built the way `WorkflowExecutionServiceSpec`/`WorkflowServiceSpec` do: + * construction does no external work, so the coordinator config and result service can be null. + */ + private def newExecution(): WorkflowExecutionService = + new WorkflowExecutionService( + null, + new WorkflowContext(), + null, + executeRequest, + new ExecutionStateStore(), + (_: Throwable) => (), + None, + new URI("vfs:///test") + ) + + /** A session state holding `workflow`, with `execution` published into it when given. */ + private def attach( + session: Session, + access: PrivilegeEnum, + workflow: TestWorkflowService, + execution: Option[WorkflowExecutionService] = None + ): SessionState = { + val state = registerState(session, access) + state.subscribe(workflow) + execution.foreach(workflow.executionService.onNext) + state } private def executeRequest: WorkflowExecuteRequest = @@ -184,6 +353,14 @@ class WorkflowWebsocketResourceSpec private def frameOf(request: TexeraWebSocketRequest): String = objectMapper.writeValueAsString(request) + /** Shared by the "forwarded" and "no workflow attached" halves of each guarded arm. */ + private val modifyLogicFrame = + """{"type":"ModifyLogicRequest","operator":{"operatorType":"Limit","limit":7}}""" + + private val paginationFrame = + """{"type":"ResultPaginationRequest","requestID":"req-77","operatorID":"op-88", + |"pageIndex":3,"pageSize":25}""".stripMargin + /** The `type` discriminator of each frame the endpoint wrote back. */ private def sentTypes(sent: ArrayBuffer[String]): Seq[String] = sent.toSeq.map(objectMapper.readTree(_).get("type").asText()) @@ -295,8 +472,267 @@ class WorkflowWebsocketResourceSpec fatalErrorMessages(sent).head should include("workflow execution is not initialized") } + it should "hand an unrecognised runtime command to the execution's websocket input" in { + // The other half of the same arm: with an execution attached the command is forwarded rather + // than rejected, and the sender's uid rides along with it. + val (session, sent) = mockSession(uid = Some(42)) + val workflow = new TestWorkflowService(9102L) + val execution = newExecution() + attach(session, PrivilegeEnum.WRITE, workflow, Some(execution)) + + val received = ArrayBuffer.empty[(TexeraWebSocketRequest, Option[Integer])] + execution.wsInput.subscribe[TexeraWebSocketRequest]((req, uid) => received += ((req, uid))) + + // A payload-carrying request, deliberately. `WorkflowPauseRequest` is a field-less case class, + // so its generated equals matches ANY instance: asserting on one would stay green with the + // deserialized request dropped and a fresh constant forwarded in its place. + resource.myOnMsg(session, """{"type":"RetryRequest","workers":["worker-a","worker-b"]}""") + + received should have size 1 + received.head._1 shouldBe RetryRequest(Seq("worker-a", "worker-b")) + // Some(42), not None: None is what a missing user-properties entry yields anyway, so asserting + // it would be green with the whole uid lookup deleted. + received.head._2 shouldBe Some(Integer.valueOf(42)) + // Forwarding is the whole handling — nothing goes back to the socket, and no exception escapes. + sent shouldBe empty + } + + // -- reconfiguration ---------------------------------------------------------- + + it should "forward a ModifyLogicRequest to the execution's reconfiguration service" in { + val (session, sent) = mockSession(uid = Some(42)) + val workflow = new TestWorkflowService(9103L) + val execution = newExecution() + val reconfiguration = new CapturingReconfigurationService(execution.executionStateStore) + execution.executionReconfigurationService = reconfiguration + attach(session, PrivilegeEnum.WRITE, workflow, Some(execution)) + + resource.myOnMsg(session, modifyLogicFrame) + + // The deserialized operator has to arrive intact: asserting only that SOME ModifyLogicResponse + // came back would stay green with a different (or default) request forwarded. + val captured = reconfiguration.captured.getOrElse(fail("no request reached the service")) + captured.operator shouldBe a[LimitOpDesc] + captured.operator.asInstanceOf[LimitOpDesc].limit shouldBe 7 + + sentTypes(sent) shouldBe Seq("ModifyLogicResponse") + val response = objectMapper.readTree(sent.head) + // Both values were computed by the stub FROM the request it received, so this pins the whole + // round trip rather than the shape of a canned reply. + response.get("opId").asText() shouldBe captured.operator.operatorIdentifier.id + response.get("errorMessage").asText() shouldBe "7" + } + + it should "ignore a ModifyLogicRequest that arrives before any workflow is attached" in { + // The `if (workflowStateOpt.isDefined)` guard, observed false. Without this case the guard could + // be replaced by `if (true)` and nothing would notice: the `.get` on the empty Option would + // throw NoSuchElementException, the catch-all would turn it into a WorkflowErrorEvent, and the + // rethrow would escape — so both assertions here have teeth. + // + // NOT pinned, and it cannot be: rewriting the guard as `if (executionStateOpt.isDefined)` is + // indistinguishable from the current one to every test in this suite. The only input that + // separates them is a workflow with no execution, where today's guard NPEs on + // `executionService.getValue` and the rewrite quietly does nothing. Pinning either answer would + // cement one of them, and the rewrite is arguably the fix. + val (session, sent) = mockSession(uid = Some(42)) + registerState(session, PrivilegeEnum.WRITE) + + noException should be thrownBy resource.myOnMsg(session, modifyLogicFrame) + + sent shouldBe empty + } + + // -- pagination --------------------------------------------------------------- + + it should "answer a pagination request with what the result service returns" in { + val (session, sent) = mockSession() + val workflow = new TestWorkflowService(9104L) + attach(session, PrivilegeEnum.WRITE, workflow) + + resource.myOnMsg(session, paginationFrame) + + // Four distinct field values, so the request cannot have been rebuilt from defaults on the way + // in; the endpoint has to pass the deserialized one straight through. + val captured = + workflow.results.captured.getOrElse(fail("no request reached the result service")) + captured.requestID shouldBe "req-77" + captured.operatorID shouldBe "op-88" + captured.pageIndex shouldBe 3 + captured.pageSize shouldBe 25 + + sentTypes(sent) shouldBe Seq("PaginatedResultEvent") + val event = objectMapper.readTree(sent.head) + event.get("requestID").asText() shouldBe "req-77" + event.get("operatorID").asText() shouldBe "op-88" + event.get("pageIndex").asInt() shouldBe 3 + // The other direction, and the half the three fields above cannot see: those are all + // reconstructible from the request, and production has a `PaginatedResultEvent(req, ...)` + // companion that does exactly that. The table and schema exist only in the event the service + // returned, so they are what proves the endpoint forwards the RETURN VALUE. + event.get("table").get(0).get("c").asInt() shouldBe 5 + event.get("schema").get(0).get("attributeName").asText() shouldBe "c" + } + + it should "ignore a pagination request that arrives before any workflow is attached" in { + // `workflowStateOpt.foreach(...)`, observed empty. Replacing that with + // `workflowStateOpt.get.resultService...` throws NoSuchElementException here, which the + // catch-all turns into a WorkflowErrorEvent and then rethrows — so both assertions bite. + val (session, sent) = mockSession() + registerState(session, PrivilegeEnum.WRITE) + + noException should be thrownBy resource.myOnMsg(session, paginationFrame) + + sent shouldBe empty + } + + // -- the error mapper's metadata-store arm ------------------------------------ + + it should "record a failure in the execution's metadata store, replacing the stale compilation error" in { + // With an execution attached the mapper writes into its metadata store instead of the socket, + // and drops any earlier COMPILATION_ERROR so the panel shows one compilation failure at a time. + val (session, sent) = mockSession(uid = Some(42)) + val workflow = new TestWorkflowService(9105L) + val execution = newExecution() + attach(session, PrivilegeEnum.READ, workflow, Some(execution)) + + val store = execution.executionStateStore.metadataStore + // Two seeded errors of two different types. With only the new error in play the filter would + // have nothing to remove, so deleting it entirely would leave the test green. + store.updateState( + _.addFatalErrors( + WorkflowFatalError( + COMPILATION_ERROR, + Timestamp(Instant.now), + "stale compilation", + "", + "op-a" + ), + WorkflowFatalError(EXECUTION_FAILURE, Timestamp(Instant.now), "runtime failure", "", "op-b") + ) + ) + + // READ access, so the write-access gate throws inside the try. + intercept[IllegalStateException] { + resource.myOnMsg(session, frameOf(executeRequest)) + } + + val errors = store.getState.fatalErrors + // The unrelated execution failure survives; only the compilation error is superseded. + errors.map(_.message) should contain("runtime failure") + errors.map(_.message) should not contain "stale compilation" + val compilationErrors = errors.filter(_.`type` == COMPILATION_ERROR) + compilationErrors should have size 1 + compilationErrors.head.message should include("write access") + // `details` is what the frontend's error panel shows a developer, and it has to be the stack + // trace, not the message again: naming the throwing frame is what separates + // `getStackTraceWithAllCauses(err)` from `err.toString` / `err.getMessage`, both of which would + // satisfy any "is non-empty" check. + compilationErrors.head.details should include("WorkflowWebsocketResource.myOnMsg") + // A real clock reading, not `Timestamp.defaultInstance` — the panel orders errors by this. + compilationErrors.head.timestamp.seconds should be > 0L + + // Nothing reaches the socket: this is the arm that replaces the WorkflowErrorEvent send. + sent shouldBe empty + } + + it should "route a failure raised while the execution is being created to the socket" in { + // `executionStateOpt` is captured BEFORE the try, so it names the execution that existed when + // the frame arrived. An execution created DURING the request therefore does not receive the + // error — it goes to the socket instead. Recomputing that capture inside the catch (it is used + // nowhere else) would reverse the routing, and no other case in this suite can see the + // difference, because every other one either has an execution throughout or never gains one. + val (session, sent) = mockSession(uid = Some(42)) + val workflow = new TestWorkflowService(9106L) + val execution = newExecution() + workflow.publishThenFail = Some(execution) + // Attached with NO execution: the capture is None. + attach(session, PrivilegeEnum.WRITE, workflow) + + intercept[IllegalStateException] { + resource.myOnMsg(session, frameOf(executeRequest)) + } + + // The execution exists by the time the handler fails ... + workflow.executionService.getValue shouldBe execution + // ... and is nonetheless not where the error was recorded. + execution.executionStateStore.metadataStore.getState.fatalErrors shouldBe empty + sentTypes(sent) shouldBe Seq("WorkflowStateEvent", "WorkflowErrorEvent") + fatalErrorMessages(sent).head should include("failed while starting the execution") + } + // -- session lifecycle -------------------------------------------------------- + "myOnOpen" should "bind the session to the requested workflow, computing unit and privilege" in { + // Two rows, because one cannot tell a real parse from a constant: the ids differ from each + // other and from 1, and the privilege differs between rows. + Seq((970101L, 7, PrivilegeEnum.READ), (970102L, 13, PrivilegeEnum.WRITE)).foreach { + case (wid, cuid, access) => + withClue(s"wid=$wid cuid=$cuid access=$access: ") { + val (session, sent) = mockSession(cuAccess = access.name(), wid = wid, cuid = cuid) + // The service this call has to JOIN. Created up front so the assertion below can be about + // instance identity rather than field values. + val shared = holdWorkflowService(wid, cuid) + // Registered before the call so afterEach still cleans up if myOnOpen throws. + registeredSessions += session.getId + + // A sentinel node count, so the event's payload is pinned to the global the endpoint is + // supposed to read rather than to 0, which is that global's initializer. Restored in a + // finally, the way ClusterListenerSpec handles the same var — amber serialises its suites + // (amber/build.sbt: `Tags.limit(Tags.Test, 1)`), so the window is this call alone. + val previousNodeCount = ClusterListener.numWorkerNodesInCluster + ClusterListener.numWorkerNodesInCluster = 4242 + try resource.myOnOpen(session, null) + finally ClusterListener.numWorkerNodesInCluster = previousNodeCount + + val state = SessionState.getState(session.getId) + // The same INSTANCE, not merely the same ids. `WorkflowService.getOrCreate` is what makes + // one workflow shared by every client that opens it; `new WorkflowService(...)` in its + // place would satisfy every field assertion below while silently handing each websocket a + // private workflow of its own — i.e. deleting multi-user collaboration. + state.getCurrentWorkflowState.getOrElse( + fail("no workflow bound") + ) should be theSameInstanceAs shared + state.getCurrentWorkflowState.map(s => (s.workflowId.id, s.computingUnitId)) shouldBe + Some((wid, cuid)) + // READ/WRITE, never NONE: NONE is SessionState's own field default, so feeding it would + // pass with the privilege never being set at all. + state.getUserComputingUnitAccess shouldBe access + + // Ordered, not merely present. The state event is the "hack to refresh frontend run + // button state" and has to go out before the workflow subscription is established. + sentTypes(sent) shouldBe Seq("WorkflowStateEvent", "ClusterStatusUpdateEvent") + objectMapper.readTree(sent.head).get("state").asText() shouldBe "Uninitialized" + objectMapper.readTree(sent(1)).get("numWorkers").asInt() shouldBe 4242 + } + } + } + + it should "join a workflow that already exists, ignoring the computing unit it was opened with" in { + // The cached arm of `getOrCreate`, which the rows above never take. `workflowServiceMapping` is + // keyed on `mkWorkflowStateId(workflowId)` — the workflow id ALONE — so `computingUnitId` is + // used only when the entry is constructed and is silently dropped for every later opener. + // + // This is characterization, not endorsement: a second websocket asking for the same workflow on + // a DIFFERENT computing unit is bound to the first one's, with nothing reported. Recorded here + // because it is the behaviour the endpoint actually has, and because the test above cannot see + // it — both of its rows construct. + val existing = holdWorkflowService(wid = 970103L, cuid = 5) + val (session, _) = mockSession(cuAccess = PrivilegeEnum.WRITE.name(), wid = 970103L, cuid = 99) + registeredSessions += session.getId + + resource.myOnOpen(session, null) + + val bound = SessionState + .getState(session.getId) + .getCurrentWorkflowState + .getOrElse(fail("no workflow bound")) + // Identity is the assertion, deliberately: it is what pins the cached arm (keying `getOrCreate` + // on wid+cuid would construct a second service and fail this). The dropped computing unit id is + // described above but NOT asserted -- pinning it would turn a defect into a contract and make + // this spec an obstacle to fixing it. + bound should be theSameInstanceAs existing + } + "myOnClose" should "drop the state registered under that session id" in { val (session, _) = mockSession() val state = registerState(session, PrivilegeEnum.WRITE)
