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-6438-8df38afe777b1240c5c9263fb16862ccb199f441 in repository https://gitbox.apache.org/repos/asf/texera.git
commit e08b7f336af68ce1bcaca012589beb41c7aed132 Author: Prateek Ganigi <[email protected]> AuthorDate: Thu Jul 16 13:47:58 2026 -0700 test(workflow-operator): add unit tests for MediaGenCodegen (#6438) ### What changes were proposed in this PR? Adds dedicated unit test coverage for MediaGenCodegen, the TaskCodegen object that generates the Python payload/parse snippets for Hugging Face media-generation operators (text-to-image and text-to-video). This class previously had no direct test coverage. A new spec file, [MediaGenCodegenSpec.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 9 tests covering: - task equals "text-to-image" - tasks equals exactly Set("text-to-image", "text-to-video") - payloadPython emits the static {"inputs": prompt_value} payload - parsePython branches on both media tasks, routes URL responses through the shared _url_to_data_url helper, converts OpenAI b64_json payloads into a data:image/png;base64, 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 #6252 ### How was this PR tested? Added 9 unit tests in the new spec file, all passing via sbt: sbt "WorkflowOperator/testOnly org.apache.texera.amber.operator.huggingFace.codegen.MediaGenCodegenSpec" [info] Tests: succeeded 9, 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/MediaGenCodegenSpec.scala | 140 +++++++++++++++++++++ 1 file changed, 140 insertions(+) diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/MediaGenCodegenSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/MediaGenCodegenSpec.scala new file mode 100644 index 0000000000..02a4c2d607 --- /dev/null +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/codegen/MediaGenCodegenSpec.scala @@ -0,0 +1,140 @@ +/* + * 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 MediaGenCodegenSpec extends AnyFlatSpec with Matchers { + + private def makeCtx( + hfApiToken: EncodableString = "token", + modelId: EncodableString = "stabilityai/stable-diffusion-2", + promptColumn: EncodableString = "prompt", + resultColumn: EncodableString = "hf_response", + task: EncodableString = "text-to-image", + systemPrompt: EncodableString = "You are a helpful assistant.", + safeMaxTokens: Int = 256, + safeTemp: Double = 0.7 + ): CodegenContext = + CodegenContext( + hfApiToken = hfApiToken, + modelId = modelId, + promptColumn = promptColumn, + resultColumn = resultColumn, + task = task, + systemPrompt = systemPrompt, + safeMaxTokens = safeMaxTokens, + safeTemp = safeTemp + ) + + "MediaGenCodegen.task" should "be the canonical text-to-image string" in { + MediaGenCodegen.task shouldBe "text-to-image" + } + + "MediaGenCodegen.tasks" should "cover exactly text-to-image and text-to-video" in { + MediaGenCodegen.tasks shouldBe Set("text-to-image", "text-to-video") + } + + "MediaGenCodegen.payloadPython" should "emit the static inputs payload shape" in { + val out = MediaGenCodegen.payloadPython(makeCtx()) + out should include("""payload = {"inputs": prompt_value}""") + } + + "MediaGenCodegen.parsePython" should "branch on both media tasks" in { + val out = MediaGenCodegen.parsePython(makeCtx()) + out should include("""if task == "text-to-image":""") + out should include("""elif task == "text-to-video":""") + } + + it should "route URL responses through the _url_to_data_url helper" in { + val out = MediaGenCodegen.parsePython(makeCtx()) + out should include("self._url_to_data_url(") + } + + it should "convert OpenAI b64_json payloads into a data URL" in { + val out = MediaGenCodegen.parsePython(makeCtx()) + out should include("b64_json") + out should include("data:image/png;base64,") + } + + it should "fall back to json.dumps(body) when no known shape matches" in { + val out = MediaGenCodegen.parsePython(makeCtx()) + out should include("json.dumps(body)") + } + + "MediaGenCodegen 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" + ) + val payload = MediaGenCodegen.payloadPython(ctx) + val parse = MediaGenCodegen.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" + 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" + } + + it should "produce identical output regardless of the CodegenContext contents" in { + // media generation's 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", + systemPrompt = "system-A", + safeMaxTokens = 1, + safeTemp = 0.0 + ) + val ctxB = makeCtx( + hfApiToken = "token-B", + modelId = "model-B", + promptColumn = "col-B", + resultColumn = "result-B", + systemPrompt = "system-B", + safeMaxTokens = 4096, + safeTemp = 2.0 + ) + + MediaGenCodegen.payloadPython(ctxA) shouldBe MediaGenCodegen.payloadPython(ctxB) + MediaGenCodegen.parsePython(ctxA) shouldBe MediaGenCodegen.parsePython(ctxB) + } +}
