Spaces:
Running on L40S
Running on L40S
File size: 1,455 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
import json
import pytest
from cosmos_framework.inference.dataset_samples import _normalize_caption
@pytest.mark.L0
def test_normalize_caption_string_unchanged() -> None:
sample = {"ai_caption": "Pick up the cup."}
assert _normalize_caption(sample) == "Pick up the cup."
assert sample["ai_caption"] == "Pick up the cup."
@pytest.mark.L0
def test_normalize_caption_dict_serialized_and_written_back() -> None:
caption = {
"cinematography": {"framing": "third-person view"},
"actions": [{"time": "0:00-0:02", "description": "Open the drawer."}],
"duration": "2s",
"fps": 8.0,
"resolution": {"H": 192, "W": 320},
"aspect_ratio": "16,9",
}
sample = {"ai_caption": caption}
result = _normalize_caption(sample)
assert json.loads(result) == caption
assert sample["ai_caption"] == result
@pytest.mark.L0
def test_normalize_caption_missing_returns_empty() -> None:
assert _normalize_caption({}) == ""
@pytest.mark.L0
def test_normalize_caption_raises_on_non_str_non_dict() -> None:
with pytest.raises(TypeError, match="ai_caption must be str or dict"):
_normalize_caption({"ai_caption": None})
with pytest.raises(TypeError, match="ai_caption must be str or dict"):
_normalize_caption({"ai_caption": 42})
|