[AZ-171] Add TensorRT tests, AC coverage gate in implement skill, optimize test infrastructure

- Add TensorRT export tests with graceful skip when no GPU available
- Add AC test coverage verification step (Step 8) to implement skill
- Add test coverage gap analysis to new-task skill
- Move exported_models fixture to conftest.py as session-scoped (shared across modules)
- Reorder tests: e2e training runs first so images/labels are available for all tests
- Consolidate teardown into single session-level cleanup in conftest.py
- Fix infrastructure tests to count files dynamically instead of hardcoded 20

Made-with: Cursor

tests/conftest.py

@@ -1,5 +1,6 @@
 import csv
 import shutil
+import sys
 from pathlib import Path
 
 import pytest
@@ -12,12 +13,21 @@ _DATASET_LABELS = _TEST_ROOT / "data" / "labels"
 _ONNX_MODEL = _PROJECT_ROOT / "_docs/00_problem/input_data/azaion.onnx"
 _CLASSES_JSON = _PROJECT_ROOT / "src" / "classes.json"
 _CONFIG_TEST = _PROJECT_ROOT / "config.test.yaml"
+_MODELS_DIR = _TEST_ROOT / "models"
 
 collect_ignore = ["security_test.py", "imagelabel_visualize_test.py"]
 
+_E2E_MODULE = "test_training_e2e"
+
 _test_results = []
 
 
+def pytest_collection_modifyitems(items):
+    e2e = [i for i in items if _E2E_MODULE in i.nodeid]
+    rest = [i for i in items if _E2E_MODULE not in i.nodeid]
+    items[:] = e2e + rest
+
+
 @pytest.hookimpl(tryfirst=True, hookwrapper=True)
 def pytest_runtest_makereport(item, call):
     outcome = yield
@@ -32,16 +42,21 @@ def pytest_runtest_makereport(item, call):
 
 
 def pytest_sessionfinish(session, exitstatus):
-    if not _test_results:
-        return
-    results_dir = Path(__file__).resolve().parent / "test-results"
-    results_dir.mkdir(exist_ok=True)
+    if _test_results:
+        results_dir = Path(__file__).resolve().parent / "test-results"
+        results_dir.mkdir(exist_ok=True)
 
-    with open(results_dir / "test-results.csv", "w", newline="", encoding="utf-8") as f:
-        writer = csv.writer(f)
-        writer.writerow(["module", "test", "result", "duration_s"])
-        for r in _test_results:
-            writer.writerow([r["module"], r["name"], r["result"], f"{r['duration']:.3f}"])
+        with open(results_dir / "test-results.csv", "w", newline="", encoding="utf-8") as f:
+            writer = csv.writer(f)
+            writer.writerow(["module", "test", "result", "duration_s"])
+            for r in _test_results:
+                writer.writerow([r["module"], r["name"], r["result"], f"{r['duration']:.3f}"])
+
+    import constants as c
+    test_config = c.Config.from_yaml(str(_CONFIG_TEST), root=str(_TEST_ROOT))
+    for d in (_DATASET_IMAGES, _DATASET_LABELS, test_config.datasets_dir,
+              test_config.corrupted_dir, str(_MODELS_DIR)):
+        shutil.rmtree(str(d), ignore_errors=True)
 
 
 def apply_constants_patch(monkeypatch, base: Path):
@@ -157,3 +172,33 @@ def empty_label(tmp_path):
     p.parent.mkdir(parents=True, exist_ok=True)
     p.write_text("", encoding="utf-8")
     return p
+
+
+@pytest.fixture(scope="session")
+def exported_models():
+    from ultralytics import YOLO
+    import constants as c
+    import exports as exports_mod
+
+    _MODELS_DIR.mkdir(parents=True, exist_ok=True)
+
+    pt_path = str(_MODELS_DIR / "test.pt")
+    YOLO("yolo11n.pt").save(pt_path)
+
+    old_config = c.config
+    c.config = c.Config.from_yaml(str(_CONFIG_TEST), root=str(_TEST_ROOT))
+    imgsz = c.config.export.onnx_imgsz
+
+    exports_mod.export_onnx(pt_path)
+    if sys.platform == "darwin":
+        exports_mod.export_coreml(pt_path)
+
+    c.config = old_config
+
+    onnx_files = list(_MODELS_DIR.glob("test*.onnx"))
+    return {
+        "onnx": str(onnx_files[0]) if onnx_files else None,
+        "model_dir": _MODELS_DIR,
+        "pt_path": pt_path,
+        "imgsz": imgsz,
+    }

tests/test_export.py

@@ -5,41 +5,23 @@ import cv2
 import numpy as np
 import onnxruntime as ort
 import pytest
+import torch
 from ultralytics import YOLO
 
 import constants as c
 import exports as exports_mod
 
+_HAS_TENSORRT = torch.cuda.is_available()
+try:
+    import tensorrt
+except ImportError:
+    _HAS_TENSORRT = False
+
 _TESTS_DIR = Path(__file__).resolve().parent
 _CONFIG_TEST = _TESTS_DIR.parent / "config.test.yaml"
 _DATASET_IMAGES = _TESTS_DIR / "root" / "data" / "images"
 
 
-@pytest.fixture(scope="module")
-def exported_models(tmp_path_factory):
-    # Arrange
-    tmp = tmp_path_factory.mktemp("export")
-    model_dir = tmp / "models"
-    model_dir.mkdir()
-
-    pt_path = str(model_dir / "test.pt")
-    YOLO("yolo11n.pt").save(pt_path)
-
-    old_config = c.config
-    c.config = c.Config.from_yaml(str(_CONFIG_TEST), root=str(tmp))
-
-    # Act
-    exports_mod.export_onnx(pt_path)
-    exports_mod.export_coreml(pt_path)
-
-    yield {
-        "onnx": str(next(model_dir.glob("*.onnx"))),
-        "model_dir": model_dir,
-    }
-
-    c.config = old_config
-
-
 class TestOnnxExport:
     def test_onnx_file_created(self, exported_models):
         # Assert
@@ -59,7 +41,7 @@ class TestOnnxExport:
         # Arrange
         session = ort.InferenceSession(exported_models["onnx"], providers=["CPUExecutionProvider"])
         meta = session.get_inputs()[0]
-        imgsz = c.config.export.onnx_imgsz
+        imgsz = exported_models["imgsz"]
         imgs = sorted(_DATASET_IMAGES.glob("*.jpg"))
         if not imgs:
             pytest.skip("no test images")
@@ -77,7 +59,7 @@ class TestOnnxExport:
         # Arrange
         session = ort.InferenceSession(exported_models["onnx"], providers=["CPUExecutionProvider"])
         meta = session.get_inputs()[0]
-        imgsz = c.config.export.onnx_imgsz
+        imgsz = exported_models["imgsz"]
         imgs = sorted(_DATASET_IMAGES.glob("*.jpg"))
         if not imgs:
             pytest.skip("no test images")
@@ -93,6 +75,82 @@ class TestOnnxExport:
         assert out[0].shape[0] == 4
 
 
+@pytest.mark.skipif(not _HAS_TENSORRT, reason="TensorRT requires NVIDIA GPU and tensorrt package")
+class TestTensorrtExport:
+    @pytest.fixture(scope="class")
+    def tensorrt_model(self, exported_models):
+        # Arrange
+        model_dir = exported_models["model_dir"]
+        pt_path = exported_models["pt_path"]
+        old_config = c.config
+        c.config = c.Config.from_yaml(str(_CONFIG_TEST), root=str(model_dir.parent))
+
+        # Act
+        exports_mod.export_tensorrt(pt_path)
+
+        c.config = old_config
+        engines = list(model_dir.glob("*.engine"))
+        yield {
+            "engine": str(engines[0]) if engines else None,
+            "model_dir": model_dir,
+            "imgsz": exported_models["imgsz"],
+        }
+
+        for e in model_dir.glob("*.engine"):
+            e.unlink(missing_ok=True)
+
+    def test_tensorrt_engine_created(self, tensorrt_model):
+        # Assert
+        assert tensorrt_model["engine"] is not None
+        p = Path(tensorrt_model["engine"])
+        assert p.exists()
+        assert p.stat().st_size > 0
+
+    def test_tensorrt_inference_batch_1(self, tensorrt_model):
+        # Arrange
+        assert tensorrt_model["engine"] is not None
+        imgs = sorted(_DATASET_IMAGES.glob("*.jpg"))
+        if not imgs:
+            pytest.skip("no test images")
+        model = YOLO(tensorrt_model["engine"])
+
+        # Act
+        results = model.predict(source=str(imgs[0]), imgsz=tensorrt_model["imgsz"], verbose=False)
+
+        # Assert
+        assert len(results) == 1
+        assert results[0].boxes is not None
+
+    def test_tensorrt_inference_batch_multiple(self, tensorrt_model):
+        # Arrange
+        assert tensorrt_model["engine"] is not None
+        imgs = sorted(_DATASET_IMAGES.glob("*.jpg"))
+        if len(imgs) < 4:
+            pytest.skip("need at least 4 test images")
+        model = YOLO(tensorrt_model["engine"])
+
+        # Act
+        results = model.predict(source=[str(p) for p in imgs[:4]], imgsz=tensorrt_model["imgsz"], verbose=False)
+
+        # Assert
+        assert len(results) == 4
+
+    def test_tensorrt_inference_batch_max(self, tensorrt_model):
+        # Arrange
+        assert tensorrt_model["engine"] is not None
+        imgs = sorted(_DATASET_IMAGES.glob("*.jpg"))
+        if not imgs:
+            pytest.skip("no test images")
+        model = YOLO(tensorrt_model["engine"])
+        sources = [str(imgs[0])] * 8
+
+        # Act
+        results = model.predict(source=sources, imgsz=tensorrt_model["imgsz"], verbose=False)
+
+        # Assert
+        assert len(results) == 8
+
+
 @pytest.mark.skipif(sys.platform != "darwin", reason="CoreML requires macOS")
 class TestCoremlExport:
     def test_coreml_package_created(self, exported_models):
@@ -117,7 +175,7 @@ class TestCoremlExport:
         model = YOLO(str(pkgs[0]))
 
         # Act
-        results = model.predict(source=str(imgs[0]), imgsz=c.config.export.onnx_imgsz, verbose=False)
+        results = model.predict(source=str(imgs[0]), imgsz=exported_models["imgsz"], verbose=False)
 
         # Assert
         assert len(results) == 1

tests/test_infrastructure.py

@@ -3,12 +3,14 @@ import constants as c
 
 def test_fixture_images_dir_has_jpegs(fixture_images_dir):
     jpgs = list(fixture_images_dir.glob("*.jpg"))
-    assert len(jpgs) == 20
+    assert len(jpgs) > 0
 
 
-def test_fixture_labels_dir_has_yolo_labels(fixture_labels_dir):
+def test_fixture_labels_dir_has_yolo_labels(fixture_labels_dir, fixture_images_dir):
     txts = list(fixture_labels_dir.glob("*.txt"))
-    assert len(txts) == 20
+    jpgs = list(fixture_images_dir.glob("*.jpg"))
+    assert len(txts) > 0
+    assert len(txts) == len(jpgs)
 
 
 def test_fixture_onnx_model_bytes(fixture_onnx_model):

tests/test_training_e2e.py

@@ -62,11 +62,6 @@ def e2e_result():
         "linked_count": linked_count,
     }
 
-    shutil.rmtree(str(dst_images), ignore_errors=True)
-    shutil.rmtree(str(dst_labels), ignore_errors=True)
-    shutil.rmtree(c.config.datasets_dir, ignore_errors=True)
-    shutil.rmtree(c.config.models_dir, ignore_errors=True)
-    shutil.rmtree(c.config.corrupted_dir, ignore_errors=True)
     c.config = old_config