From 2b95db5bcb38dbcccdeb5c5dba2c33b88cfd6310 Mon Sep 17 00:00:00 2001
From: Pmobilee <muldamion@gmail.com>
Date: Mon, 18 May 2026 16:02:25 +0900
Subject: [PATCH] Update tests/test_frame_selector.py for temporal frame
 ranking

---
 tests/test_frame_selector.py | 45 ++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/tests/test_frame_selector.py b/tests/test_frame_selector.py
index 6cd1404..aba0419 100644
--- a/tests/test_frame_selector.py
+++ b/tests/test_frame_selector.py
@@ -13,6 +13,7 @@ SPEC.loader.exec_module(frame_selector)
 
 AIA_SelectVideoFrame = frame_selector.AIA_SelectVideoFrame
 _normalize_selector = frame_selector._normalize_selector
+_selector_rank = frame_selector._selector_rank
 
 
 def _checker(size=64, block=8):
@@ -42,6 +43,27 @@ def test_selector_aliases():
     assert _normalize_selector("static pose") == "static"
     assert _normalize_selector("strong lighting flicker") == "lighting"
     assert _normalize_selector("whatever") == "best"
+    assert _selector_rank("moving") == 1
+    assert _selector_rank("moving rank 2") == 2
+    assert _selector_rank("lighting #3") == 3
+
+
+def test_ranked_indices_spread_candidates_apart():
+    metrics = {
+        "blur": np.asarray([0, 10, 100, 99, 20, 10, 0, 94, 0, 93], dtype=np.float32),
+        "motion": np.zeros(10, dtype=np.float32),
+        "lighting": np.zeros(10, dtype=np.float32),
+        "face": np.ones(10, dtype=np.float32),
+    }
+
+    indices = AIA_SelectVideoFrame._ranked_indices(
+        metrics,
+        "best",
+        count=3,
+        min_separation=3,
+    )
+
+    assert indices == [2, 7, 9]
 
 
 def test_blurry_mode_selects_soft_frame(monkeypatch):
@@ -100,6 +122,29 @@ def test_no_face_mode_uses_face_scores(monkeypatch):
     assert torch.allclose(selected[0], _to_tensor([frames[2]])[0])
 
 
+def test_ranked_selector_can_pick_second_distinct_frame(monkeypatch):
+    frames = [_checker() for _ in range(40)]
+    frames[10] = np.full((64, 64, 3), 210, dtype=np.uint8)
+    frames[30] = np.full((64, 64, 3), 180, dtype=np.uint8)
+    metrics = {
+        "blur": np.asarray([1.0] * 40, dtype=np.float32),
+        "motion": np.asarray([0.0] * 40, dtype=np.float32),
+        "lighting": np.asarray([0.0] * 40, dtype=np.float32),
+        "face": np.asarray([1.0] * 40, dtype=np.float32),
+    }
+    metrics["motion"][10] = 10.0
+    metrics["motion"][30] = 9.0
+    monkeypatch.setattr(
+        AIA_SelectVideoFrame,
+        "_metrics",
+        classmethod(lambda _cls, _frames: metrics),
+    )
+
+    selected = AIA_SelectVideoFrame().select(_to_tensor(frames), "moving rank 2")[0]
+
+    assert torch.allclose(selected[0], _to_tensor([frames[30]])[0])
+
+
 def test_smoke_241_frame_batch(monkeypatch):
     frames = []
     base = _checker()