[Bugfix] Modify conftest.py set unspecified parameters (#2263)

Signed-off-by: bjf-frz <frz123db@gmail.com>
This commit is contained in:
bjf-frz authored and GitHub committed 2026-03-27 10:33:16 +00:00
1 parent c0d90b64c6
commit 637abb759f
1 file changed
+15 -13
+15 -13
View File
@@ -81,22 +81,24 @@ def assert_image_diffusion_response(
}
}
"""
extra_body = request_config.get("extra_body", {})
num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1)
assert response.images is not None, "Image response is None"
assert len(response.images) > 0, "No images in response"
assert len(response.images) == num_outputs_per_prompt, (
f"Expected {num_outputs_per_prompt} images, got {len(response.images)}"
)
extra_body = request_config.get("extra_body") or {}
num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt")
if num_outputs_per_prompt is not None:
assert len(response.images) == num_outputs_per_prompt, (
f"Expected {num_outputs_per_prompt} images, got {len(response.images)}"
)
if run_level == "advanced_model":
expected_width = extra_body["width"] # intentionally raise KeyError if missing
expected_height = extra_body["height"] # intentionally raise KeyError if missing
width = extra_body.get("width")
height = extra_body.get("height")
for img in response.images:
assert_image_valid(img, width=expected_width, height=expected_height)
if width is not None or height is not None:
for img in response.images:
assert_image_valid(img, width=width, height=height)
def assert_video_diffusion_response(
@@ -166,9 +168,9 @@ def assert_image_valid(image: Path | Image.Image, *, width: int | None = None, h
image.load()
assert image.width > 0 and image.height > 0
if width is not None:
assert image.width == width, f"Expected width={width}, got {image.width} in {image.name}"
assert image.width == width, f"Expected width={width}, got {image.width}"
if height is not None:
assert image.height == height, f"Expected height={height}, got {image.height} in {image.name}"
assert image.height == height, f"Expected height={height}, got {image.height}"
return image