Spaces:
Sleeping
Sleeping
File size: 1,496 Bytes
901e5fa | 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 46 47 48 49 50 51 52 53 54 55 | from __future__ import annotations
import sys
from pathlib import Path
import pytest
ROOT = Path(__file__).resolve().parents[1]
PKG_ROOT = ROOT / "src"
for p in (ROOT, PKG_ROOT):
sp = str(p)
if sp not in sys.path:
sys.path.insert(0, sp)
from ai_agent.agent.tools.query_utils import (
append_format_tokens,
normalize_formats,
strip_legacy_original_formats_line,
)
@pytest.mark.parametrize(
"raw, expected",
[
([" DCM ", "nii.GZ", "dcm", ""], ["dcm", "nii.gz"]),
(["png", "PNG", " jpeg "], ["png", "jpeg"]),
([], []),
],
)
def test_normalize_formats(raw, expected):
assert normalize_formats(raw) == expected
@pytest.mark.parametrize(
"query, expected_query, expected_formats",
[
(
"segment lungs\nOriginalFormats: DCM NII.GZ\nwith contrast",
"segment lungs with contrast",
["dcm", "nii.gz"],
),
("no legacy marker here", "no legacy marker here", []),
],
)
def test_strip_legacy_original_formats_line(query, expected_query, expected_formats):
cleaned_query, formats = strip_legacy_original_formats_line(query)
assert cleaned_query == expected_query
assert formats == expected_formats
def test_append_format_tokens():
out = append_format_tokens("segment lungs", ["dcm", "nii.gz"])
assert "segment lungs" in out
assert "format:DICOM" in out
assert "format:NIfTI" in out
|