(texera) 01/01: test(workflow-operator): add unit tests for QaRankingCodegen (#6655)

[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-6655-b113a7bdd968bda4bd17c564e47d7d200ba02898
in repository https://gitbox.apache.org/repos/asf/texera.git

commit 2a9c1e3022107a48a6470caf1a5630726a863059
Author: Matthew B. <mg...@uci.edu>
AuthorDate: Mon Jul 20 15:06:54 2026 -0700

    test(workflow-operator): add unit tests for QaRankingCodegen (#6655)
    
    ### What changes were proposed in this PR?
    - Add
    `common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegenSpec.scala`,
    a new ScalaTest spec for QaRankingCodegen, which previously had no
    dedicated unit tests.
    - Cover the task and tasks mapping and primary-task membership.
    - Cover per-task payload and parse branch selection and the expected
    structural markers in the generated Python.
    - Confirm snippet determinism (byte-identical output across two
    unrelated CodegenContexts) and no leakage of ctx string values.
    ### Any related issues, documentation, discussions?
    Closes: #6654
    ### How was this PR tested?
    - Run: `sbt "WorkflowOperator/testOnly *QaRankingCodegenSpec"`, expect
    all 13 tests passing.
    - Test-only change; no production code is modified.
    ### Was this PR authored or co-authored using generative AI tooling?
    Co-authored with Claude Opus 4.8 in compliance with ASF
---
 .../huggingFace/codegen/QaRankingCodegenSpec.scala | 197 +++++++++++++++++++++
 1 file changed, 197 insertions(+)

diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegenSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegenSpec.scala
new file mode 100644
index 0000000000..0e21824b88
--- /dev/null
+++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/QaRankingCodegenSpec.scala
@@ -0,0 +1,197 @@
+/*
+ * 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.operator.huggingFace.codegen
+
+import org.apache.texera.amber.pybuilder.PyStringTypes.EncodableString
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+class QaRankingCodegenSpec extends AnyFlatSpec with Matchers {
+
+  private def makeCtx(
+      hfApiToken: EncodableString = "token",
+      modelId: EncodableString = "deepset/roberta-base-squad2",
+      promptColumn: EncodableString = "prompt",
+      resultColumn: EncodableString = "hf_response",
+      task: EncodableString = "question-answering",
+      systemPrompt: EncodableString = "You are a helpful assistant.",
+      safeMaxTokens: Int = 256,
+      safeTemp: Double = 0.7,
+      contextColumn: EncodableString = "context",
+      candidateLabels: EncodableString = "positive,negative",
+      sentencesColumn: EncodableString = "sentences"
+  ): CodegenContext =
+    CodegenContext(
+      hfApiToken = hfApiToken,
+      modelId = modelId,
+      promptColumn = promptColumn,
+      resultColumn = resultColumn,
+      task = task,
+      systemPrompt = systemPrompt,
+      safeMaxTokens = safeMaxTokens,
+      safeTemp = safeTemp,
+      contextColumn = contextColumn,
+      candidateLabels = candidateLabels,
+      sentencesColumn = sentencesColumn
+    )
+
+  "QaRankingCodegen.task" should "be the canonical question-answering string" in {
+    QaRankingCodegen.task shouldBe "question-answering"
+  }
+
+  "QaRankingCodegen.tasks" should "cover exactly the five QA/ranking task families" in {
+    QaRankingCodegen.tasks shouldBe Set(
+      "question-answering",
+      "table-question-answering",
+      "zero-shot-classification",
+      "sentence-similarity",
+      "text-ranking"
+    )
+  }
+
+  it should "include its primary task among the handled tasks" in {
+    QaRankingCodegen.tasks should contain(QaRankingCodegen.task)
+  }
+
+  "QaRankingCodegen.payloadPython" should "branch on each of the five tasks and an else fallback" in {
+    val out = QaRankingCodegen.payloadPython(makeCtx())
+    out should include("""if task == "question-answering":""")
+    out should include("""elif task == "table-question-answering":""")
+    out should include("""elif task == "zero-shot-classification":""")
+    out should include("""elif task == "sentence-similarity":""")
+    out should include("""elif task == "text-ranking":""")
+    out should include("else:")
+  }
+
+  it should "build the question-answering payload from prompt_value and the context column" in {
+    val out = QaRankingCodegen.payloadPython(makeCtx())
+    out should include("self.CONTEXT_COLUMN")
+    out should include("""payload = {"inputs": {"question": prompt_value, "context": ctx_val}}""")
+  }
+
+  it should "route table-question-answering through query and table_dict" in {
+    val out = QaRankingCodegen.payloadPython(makeCtx())
+    out should include("""payload = {"inputs": {"query": prompt_value, "table": table_dict}}""")
+  }
+
+  it should "derive zero-shot candidate labels from the candidate-labels attribute" in {
+    val out = QaRankingCodegen.payloadPython(makeCtx())
+    out should include("self.CANDIDATE_LABELS")
+    out should include("candidate_labels")
+  }
+
+  it should "split the sentences column for sentence-similarity and text-ranking" in {
+    val out = QaRankingCodegen.payloadPython(makeCtx())
+    out should include("self.SENTENCES_COLUMN")
+    out should include("source_sentence")
+    out should include(""""query": prompt_value""")
+    out should include(""""texts": sentences_list""")
+  }
+
+  it should "fall back to shipping the raw prompt as inputs" in {
+    val out = QaRankingCodegen.payloadPython(makeCtx())
+    out should include("""payload = {"inputs": prompt_value}""")
+  }
+
+  "QaRankingCodegen.parsePython" should "extract the answer field for both QA variants" in {
+    val out = QaRankingCodegen.parsePython(makeCtx())
+    out should include("""if task == "question-answering":""")
+    out should include("""elif task == "table-question-answering":""")
+    out should include("""body.get("answer"""")
+  }
+
+  it should "return the raw JSON body for the ranking-style tasks" in {
+    val out = QaRankingCodegen.parsePython(makeCtx())
+    out should include(
+      """elif task in ("zero-shot-classification", "sentence-similarity", "text-ranking"):"""
+    )
+    out should include("return json.dumps(body)")
+  }
+
+  "QaRankingCodegen snippets" should "never inline raw CodegenContext string values" in {
+    // The snippets reference only self.* attributes; the base class decodes
+    // user-supplied strings safely at runtime. Sentinel values are distinctive
+    // and non-overlapping with the static template text.
+    val ctx = makeCtx(
+      hfApiToken = "MARKER_TOKEN_zXyq42",
+      modelId = "MARKER_MODEL_zXyq42",
+      promptColumn = "MARKER_PROMPT_zXyq42",
+      resultColumn = "MARKER_RESULT_zXyq42",
+      task = "MARKER_TASK_zXyq42",
+      systemPrompt = "MARKER_SYSTEM_zXyq42",
+      contextColumn = "MARKER_CONTEXT_zXyq42",
+      candidateLabels = "MARKER_LABELS_zXyq42",
+      sentencesColumn = "MARKER_SENTENCES_zXyq42"
+    )
+    val payload = QaRankingCodegen.payloadPython(ctx)
+    val parse = QaRankingCodegen.parsePython(ctx)
+
+    for (
+      marker <- Seq(
+        "MARKER_TOKEN_zXyq42",
+        "MARKER_MODEL_zXyq42",
+        "MARKER_PROMPT_zXyq42",
+        "MARKER_RESULT_zXyq42",
+        "MARKER_TASK_zXyq42",
+        "MARKER_SYSTEM_zXyq42",
+        "MARKER_CONTEXT_zXyq42",
+        "MARKER_LABELS_zXyq42",
+        "MARKER_SENTENCES_zXyq42"
+      )
+    ) {
+      payload should not include marker
+      parse should not include marker
+    }
+  }
+
+  it should "produce identical output regardless of the CodegenContext contents" in {
+    // The payload/parse snippets are static: they reference only self.*
+    // attributes, never ctx fields. Two unrelated contexts must serialise to
+    // byte-identical Python. A future refactor that accidentally consumes a
+    // ctx field will regress here.
+    val ctxA = makeCtx(
+      hfApiToken = "token-A",
+      modelId = "model-A",
+      promptColumn = "col-A",
+      resultColumn = "result-A",
+      systemPrompt = "system-A",
+      contextColumn = "ctx-A",
+      candidateLabels = "labels-A",
+      sentencesColumn = "sent-A",
+      safeMaxTokens = 1,
+      safeTemp = 0.0
+    )
+    val ctxB = makeCtx(
+      hfApiToken = "token-B",
+      modelId = "model-B",
+      promptColumn = "col-B",
+      resultColumn = "result-B",
+      systemPrompt = "system-B",
+      contextColumn = "ctx-B",
+      candidateLabels = "labels-B",
+      sentencesColumn = "sent-B",
+      safeMaxTokens = 4096,
+      safeTemp = 2.0
+    )
+
+    QaRankingCodegen.payloadPython(ctxA) shouldBe QaRankingCodegen.payloadPython(ctxB)
+    QaRankingCodegen.parsePython(ctxA) shouldBe QaRankingCodegen.parsePython(ctxB)
+  }
+}