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-6715-56d270dfc8e5797fd7b5c6ef43955cc7d7dbee2e in repository https://gitbox.apache.org/repos/asf/texera.git
commit 1c13d6b8998a1799bc205f298da9f36fb4b64b71 Author: Prateek Ganigi <[email protected]> AuthorDate: Wed Jul 22 12:56:52 2026 -0700 test(workflow-operator): add unit tests for ImageTaskCodegen (#6715) ### What changes were proposed in this PR? Adds dedicated unit test coverage for ImageTaskCodegen, the TaskCodegen object that generates the Python payload/parse snippets for the Hugging Face image-pipeline task family (9 tasks spanning image-only, image+prompt, and zero-shot validation paths). This class previously had no direct test coverage. A new spec file, ImageTaskCodegenSpec.scala, was created in the same test package, mirroring the structure and patterns of TextGenCodegenSpec (shared makeCtx helper; assertions on snippet structure/markers rather than exact whitespace). It adds 10 tests covering: - task equals "image-classification"; tasks equals exactly the 9 image-pipeline tasks - payloadPython sends raw image bytes for image-only tasks (payload = current_image_bytes, use_raw_binary_body = True) - VQA / document-QA payloads bundle a base64 image via self._image_input_as_base64(current_image_bytes) with "question": prompt_value - Zero-shot image classification raises ValueError when fewer than 2 candidate labels are supplied parsePython extracts chat-style content, normalizes image-to-image URL responses through _url_to_data_url (including b64_json → data URL), and falls back to json.dumps(body) - Snippets never inline raw CodegenContext string values (no sentinel leakage) - Context-independence: identical output across two unrelated CodegenContext instances ### Any related issues, documentation, discussions? Closes #6250 ### How was this PR tested? Added 10 unit tests in the new spec file, all passing via sbt: sbt "WorkflowOperator/testOnly org.apache.texera.amber.operator.huggingFace.codegen.ImageTaskCodegenSpec" [info] Tests: succeeded 10, failed 0, canceled 0, ignored 0, pending 0 [info] All tests passed. ### Was this PR authored or co-authored using generative AI tooling? Co-authored with Claude Opus 4.8 in compliance with ASF. --- .../huggingFace/codegen/ImageTaskCodegenSpec.scala | 190 +++++++++++++++++++++ 1 file changed, 190 insertions(+) diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/ImageTaskCodegenSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/ImageTaskCodegenSpec.scala new file mode 100644 index 0000000000..b6e6d25aa5 --- /dev/null +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/ImageTaskCodegenSpec.scala @@ -0,0 +1,190 @@ +/* + * 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 ImageTaskCodegenSpec extends AnyFlatSpec with Matchers { + + private def makeCtx( + hfApiToken: EncodableString = "token", + modelId: EncodableString = "google/vit-base-patch16-224", + promptColumn: EncodableString = "prompt", + resultColumn: EncodableString = "hf_response", + task: EncodableString = "image-classification", + systemPrompt: EncodableString = "You are a helpful assistant.", + safeMaxTokens: Int = 256, + safeTemp: Double = 0.7, + imageInput: EncodableString = "", + inputImageColumn: EncodableString = "", + candidateLabels: EncodableString = "" + ): CodegenContext = + CodegenContext( + hfApiToken = hfApiToken, + modelId = modelId, + promptColumn = promptColumn, + resultColumn = resultColumn, + task = task, + systemPrompt = systemPrompt, + safeMaxTokens = safeMaxTokens, + safeTemp = safeTemp, + imageInput = imageInput, + inputImageColumn = inputImageColumn, + candidateLabels = candidateLabels + ) + + "ImageTaskCodegen.task" should "be the canonical image-classification string" in { + ImageTaskCodegen.task shouldBe "image-classification" + } + + "ImageTaskCodegen.tasks" should "cover exactly the nine image-pipeline tasks" in { + ImageTaskCodegen.tasks shouldBe Set( + "image-classification", + "object-detection", + "image-segmentation", + "image-to-text", + "visual-question-answering", + "document-question-answering", + "zero-shot-image-classification", + "image-text-to-text", + "image-to-image" + ) + ImageTaskCodegen.tasks should have size 9 + } + + "ImageTaskCodegen.payloadPython" should "send raw image bytes for image-only tasks" in { + val out = ImageTaskCodegen.payloadPython(makeCtx()) + out should include("if task in image_only_tasks:") + out should include("payload = current_image_bytes") + out should include("use_raw_binary_body = True") + out should include("raw_binary_headers = image_headers") + } + + it should "bundle a base64 image and question for VQA / document-QA tasks" in { + val out = ImageTaskCodegen.payloadPython(makeCtx()) + out should include( + """elif task in ("visual-question-answering", "document-question-answering"):""" + ) + out should include("self._image_input_as_base64(current_image_bytes)") + out should include(""""question": prompt_value""") + } + + it should "validate that zero-shot classification supplies at least two candidate labels" in { + val out = ImageTaskCodegen.payloadPython(makeCtx()) + out should include("""elif task == "zero-shot-image-classification":""") + out should include("if len(labels) < 2:") + out should include("raise ValueError") + out should include("candidate_labels") + } + + "ImageTaskCodegen.parsePython" should "extract chat-style content for image-text-to-text" in { + val out = ImageTaskCodegen.parsePython(makeCtx()) + out should include("choices") + out should include("message") + out should include("content") + } + + it should "normalize image-to-image URL responses through _url_to_data_url" in { + val out = ImageTaskCodegen.parsePython(makeCtx()) + out should include("self._url_to_data_url(") + out should include("data:image/png;base64,") + } + + it should "fall back to json.dumps(body) for structured tasks" in { + val out = ImageTaskCodegen.parsePython(makeCtx()) + out should include("json.dumps(body)") + } + + "ImageTaskCodegen snippets" should "never inline raw CodegenContext string values" in { + // The snippets are static and 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", + imageInput = "MARKER_IMAGE_zXyq42", + inputImageColumn = "MARKER_IMAGECOL_zXyq42", + candidateLabels = "MARKER_LABELS_zXyq42" + ) + val payload = ImageTaskCodegen.payloadPython(ctx) + val parse = ImageTaskCodegen.parsePython(ctx) + + payload should not include "MARKER_TOKEN_zXyq42" + payload should not include "MARKER_MODEL_zXyq42" + payload should not include "MARKER_PROMPT_zXyq42" + payload should not include "MARKER_RESULT_zXyq42" + payload should not include "MARKER_TASK_zXyq42" + payload should not include "MARKER_SYSTEM_zXyq42" + payload should not include "MARKER_IMAGE_zXyq42" + payload should not include "MARKER_IMAGECOL_zXyq42" + payload should not include "MARKER_LABELS_zXyq42" + parse should not include "MARKER_TOKEN_zXyq42" + parse should not include "MARKER_MODEL_zXyq42" + parse should not include "MARKER_PROMPT_zXyq42" + parse should not include "MARKER_RESULT_zXyq42" + parse should not include "MARKER_TASK_zXyq42" + parse should not include "MARKER_SYSTEM_zXyq42" + parse should not include "MARKER_IMAGE_zXyq42" + parse should not include "MARKER_IMAGECOL_zXyq42" + parse should not include "MARKER_LABELS_zXyq42" + } + + it should "produce identical output regardless of the CodegenContext contents" in { + // image-task payload/parse 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", + task = "image-classification", + systemPrompt = "system-A", + safeMaxTokens = 1, + safeTemp = 0.0, + imageInput = "image-A", + inputImageColumn = "image-col-A", + candidateLabels = "labels-A" + ) + val ctxB = makeCtx( + hfApiToken = "token-B", + modelId = "model-B", + promptColumn = "col-B", + resultColumn = "result-B", + task = "zero-shot-image-classification", + systemPrompt = "system-B", + safeMaxTokens = 4096, + safeTemp = 2.0, + imageInput = "image-B", + inputImageColumn = "image-col-B", + candidateLabels = "labels-B" + ) + + ImageTaskCodegen.payloadPython(ctxA) shouldBe ImageTaskCodegen.payloadPython(ctxB) + ImageTaskCodegen.parsePython(ctxA) shouldBe ImageTaskCodegen.parsePython(ctxB) + } +}
