241 lines
8.3 KiB
Python
241 lines
8.3 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from video_ai_analysis_poc.config import load_config
|
|
|
|
|
|
class ConfigTests(unittest.TestCase):
|
|
def test_loads_local_batch_yaml_and_applies_cli_overrides(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
input_dir = root / "videos"
|
|
output_dir = root / "out"
|
|
override_input = root / "override-videos"
|
|
override_output = root / "override-out"
|
|
input_dir.mkdir()
|
|
override_input.mkdir()
|
|
config_path = root / "local_batch.yaml"
|
|
config_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"input:",
|
|
f" dir: {input_dir}",
|
|
" recursive: false",
|
|
' extensions: [".mp4", ".mov"]',
|
|
"output:",
|
|
f" dir: {output_dir}",
|
|
" overwrite: false",
|
|
"ffprobe:",
|
|
" timeout_seconds: 5",
|
|
]
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
config = load_config(
|
|
config_path,
|
|
input_dir=override_input,
|
|
output_dir=override_output,
|
|
)
|
|
|
|
self.assertEqual(config["input"]["dir"], str(override_input.resolve()))
|
|
self.assertEqual(config["output"]["dir"], str(override_output.resolve()))
|
|
self.assertFalse(config["input"]["recursive"])
|
|
self.assertEqual(config["input"]["extensions"], [".mp4", ".mov"])
|
|
self.assertEqual(config["ffprobe"]["timeout_seconds"], 5)
|
|
|
|
def test_rejects_output_dir_equal_to_input_dir(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
input_dir = root / "videos"
|
|
input_dir.mkdir()
|
|
config_path = root / "local_batch.yaml"
|
|
config_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"input:",
|
|
f" dir: {input_dir}",
|
|
"output:",
|
|
f" dir: {input_dir}",
|
|
]
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
with self.assertRaisesRegex(ValueError, "output dir must not equal input dir"):
|
|
load_config(config_path)
|
|
|
|
def test_rejects_output_dir_inside_reference_project(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
input_dir = root / "videos"
|
|
input_dir.mkdir()
|
|
forbidden_output = (
|
|
Path("/Users/yoilun/AI-train/zhengxin-vlm-0413")
|
|
/ "outputs"
|
|
/ "local-batch"
|
|
)
|
|
config_path = root / "local_batch.yaml"
|
|
config_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"input:",
|
|
f" dir: {input_dir}",
|
|
"output:",
|
|
f" dir: {forbidden_output}",
|
|
]
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, "output dir must not be inside forbidden reference dir"
|
|
):
|
|
load_config(config_path)
|
|
|
|
def test_loads_nested_mapping_values(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
input_dir = root / "videos"
|
|
output_dir = root / "output"
|
|
input_dir.mkdir()
|
|
config_path = root / "local_batch.yaml"
|
|
config_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"input:",
|
|
f" dir: {input_dir}",
|
|
"output:",
|
|
f" dir: {output_dir}",
|
|
"ffmpeg:",
|
|
" codec_decoders:",
|
|
" h264: h264_cuvid",
|
|
" hevc: hevc_cuvid",
|
|
]
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
config = load_config(config_path)
|
|
|
|
self.assertEqual(
|
|
config["ffmpeg"]["codec_decoders"],
|
|
{"h264": "h264_cuvid", "hevc": "hevc_cuvid"},
|
|
)
|
|
|
|
def test_loads_prompt_block_scalar_values(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
input_dir = root / "videos"
|
|
output_dir = root / "output"
|
|
input_dir.mkdir()
|
|
config_path = root / "local_batch.yaml"
|
|
config_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"input:",
|
|
f" dir: {input_dir}",
|
|
"output:",
|
|
f" dir: {output_dir}",
|
|
"prompt:",
|
|
" system: >-",
|
|
" First instruction.",
|
|
" Second instruction.",
|
|
"",
|
|
" Final instruction.",
|
|
" user: 'Return strict JSON.'",
|
|
]
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
config = load_config(config_path)
|
|
|
|
self.assertEqual(
|
|
config["prompt"]["system"],
|
|
"First instruction.\nSecond instruction.\n\nFinal instruction.",
|
|
)
|
|
self.assertEqual(config["prompt"]["user"], "Return strict JSON.")
|
|
|
|
def test_defaults_source_mode_to_local_and_hik_cloud_section(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
input_dir = root / "videos"
|
|
output_dir = root / "output"
|
|
input_dir.mkdir()
|
|
config_path = root / "local_batch.yaml"
|
|
config_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"input:",
|
|
f" dir: {input_dir}",
|
|
"output:",
|
|
f" dir: {output_dir}",
|
|
]
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
config = load_config(config_path)
|
|
|
|
self.assertEqual(config["source"]["mode"], "local")
|
|
self.assertIn("devices", config["hik_cloud"])
|
|
self.assertIn("time_ranges", config["hik_cloud"])
|
|
|
|
def test_loads_hik_cloud_devices_and_time_ranges_as_list_of_mappings(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
input_dir = root / "videos"
|
|
output_dir = root / "output"
|
|
input_dir.mkdir()
|
|
config_path = root / "local_batch.yaml"
|
|
config_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"input:",
|
|
f" dir: {input_dir}",
|
|
"output:",
|
|
f" dir: {output_dir}",
|
|
"source:",
|
|
" mode: hik_cloud",
|
|
"hik_cloud:",
|
|
" devices:",
|
|
" - device_serial: EXAMPLE_DEVICE_SERIAL",
|
|
" channel_no: 1",
|
|
" name: front",
|
|
" time_ranges:",
|
|
' - begin: "2026-02-03 09:00:00"',
|
|
' end: "2026-02-03 10:30:00"',
|
|
]
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
config = load_config(config_path)
|
|
|
|
self.assertEqual(config["source"]["mode"], "hik_cloud")
|
|
self.assertEqual(
|
|
config["hik_cloud"]["devices"],
|
|
[
|
|
{
|
|
"device_serial": "EXAMPLE_DEVICE_SERIAL",
|
|
"channel_no": 1,
|
|
"name": "front",
|
|
}
|
|
],
|
|
)
|
|
self.assertEqual(
|
|
config["hik_cloud"]["time_ranges"],
|
|
[
|
|
{
|
|
"begin": "2026-02-03 09:00:00",
|
|
"end": "2026-02-03 10:30:00",
|
|
}
|
|
],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|