Add unit tests for key_points - skeleton & annotators (#2099)

* Refactor key point annotator tests to use reusable fixtures for test clarity and maintainability
* Apply suggestions from code review
* fix(pre_commit): 🎨 auto format pre-commit hooks
* Move `assert_image_mostly_same` helper to a shared module and update imports in tests
* Update test/key_points/test_skeletons.py

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

test/annotators/test_core.py

@@ -31,7 +31,7 @@ from supervision.annotators.utils import ColorLookup
 from supervision.detection.core import Detections
 from supervision.draw.color import Color
 from supervision.geometry.core import Position
-from test.helpers import _create_detections
+from test.helpers import _create_detections, assert_image_mostly_same
 
 
 @pytest.fixture
@@ -58,36 +58,6 @@ def gradient_image() -> np.ndarray:
     return image
 
 
-def assert_image_mostly_same(
-    original: np.ndarray, annotated: np.ndarray, similarity_threshold: float = 0.9
-) -> None:
-    """
-    Assert that the annotated image is mostly the same as the original.
-
-    Args:
-        original: Original image
-        annotated: Annotated image
-        similarity_threshold:
-          Minimum percentage of pixels that should be the same (0.0 to 1.0)
-    """
-    # Check that images have the same shape
-    assert original.shape == annotated.shape
-
-    # Calculate number of identical pixels
-    identical_pixels = np.sum(np.all(original == annotated, axis=-1))
-    total_pixels = original.shape[0] * original.shape[1]
-    similarity = identical_pixels / total_pixels
-
-    # Check that at least similarity_threshold of pixels are identical
-    assert similarity >= similarity_threshold, (
-        f"Images are only {similarity:.1%} similar, "
-        f"which is below the {similarity_threshold:.1%} threshold"
-    )
-
-    # Check that the image is not completely identical
-    assert not np.array_equal(original, annotated), "Images are completely identical"
-
-
 class TestBoxAnnotator:
     """Tests for BoxAnnotator class"""
 

test/conftest.py

@@ -0,0 +1,66 @@
+import numpy as np
+import pytest
+
+import supervision as sv
+from test.helpers import _create_key_points
+
+
+@pytest.fixture
+def scene() -> np.ndarray:
+    return np.zeros((100, 100, 3), dtype=np.uint8)
+
+
+@pytest.fixture
+def sample_key_points() -> sv.KeyPoints:
+    return _create_key_points(
+        xy=[
+            [
+                [10, 10],
+                [20, 20],
+                [30, 30],
+                [40, 40],
+                [50, 50],
+                [60, 60],
+                [70, 70],
+                [80, 80],
+                [90, 90],
+                [10, 20],
+                [20, 30],
+                [30, 40],
+                [40, 50],
+                [50, 60],
+                [60, 70],
+                [70, 80],
+                [80, 90],
+            ],
+            [
+                [10, 40],
+                [20, 50],
+                [30, 60],
+                [40, 70],
+                [50, 80],
+                [60, 90],
+                [70, 10],
+                [80, 20],
+                [90, 30],
+                [10, 50],
+                [20, 60],
+                [30, 70],
+                [40, 80],
+                [50, 90],
+                [60, 10],
+                [70, 20],
+                [80, 30],
+            ],
+        ],
+        confidence=[
+            [0.8] * 17,
+            [0.6] * 17,
+        ],
+        class_id=[0, 1],
+    )
+
+
+@pytest.fixture
+def empty_key_points() -> sv.KeyPoints:
+    return sv.KeyPoints.empty()

test/helpers.py

@@ -208,3 +208,33 @@ def assert_almost_equal(actual, expected, tolerance=1e-5):
         AssertionError: Expected 0.5, but got 0.6.
     """
     assert abs(actual - expected) < tolerance, f"Expected {expected}, but got {actual}."
+
+
+def assert_image_mostly_same(
+    original: np.ndarray, annotated: np.ndarray, similarity_threshold: float = 0.9
+) -> None:
+    """
+    Assert that the annotated image is mostly the same as the original.
+
+    Args:
+        original: Original image
+        annotated: Annotated image
+        similarity_threshold:
+          Minimum percentage of pixels that should be the same (0.0 to 1.0)
+    """
+    # Check that images have the same shape
+    assert original.shape == annotated.shape
+
+    # Calculate number of identical pixels
+    identical_pixels = np.sum(np.all(original == annotated, axis=-1))
+    total_pixels = original.shape[0] * original.shape[1]
+    similarity = identical_pixels / total_pixels
+
+    # Check that at least similarity_threshold of pixels are identical
+    assert similarity >= similarity_threshold, (
+        f"Images are only {similarity:.1%} similar, "
+        f"which is below the {similarity_threshold:.1%} threshold"
+    )
+
+    # Check that the image is not completely identical
+    assert not np.array_equal(original, annotated), "Images are completely identical"

test/key_points/test_annotators.py

@@ -0,0 +1,81 @@
+import numpy as np
+
+import supervision as sv
+from test.helpers import assert_image_mostly_same
+
+
+class TestVertexAnnotator:
+    def test_annotate_with_default_parameters(self, scene, sample_key_points):
+        """Test annotation with default parameters."""
+        annotator = sv.VertexAnnotator()
+        result = annotator.annotate(scene=scene.copy(), key_points=sample_key_points)
+
+        # Check that the scene has been modified
+        assert_image_mostly_same(
+            original=scene, annotated=result, similarity_threshold=0.8
+        )
+
+    def test_annotate_with_custom_color_and_radius(self, scene, sample_key_points):
+        """Test annotation with custom color and radius."""
+        color = sv.Color.RED
+        radius = 5
+        annotator = sv.VertexAnnotator(color=color, radius=radius)
+        result = annotator.annotate(scene=scene.copy(), key_points=sample_key_points)
+
+        # Check that the scene has been modified
+        assert_image_mostly_same(
+            original=scene, annotated=result, similarity_threshold=0.7
+        )
+
+    def test_annotate_empty_key_points(self, scene, empty_key_points):
+        """Test annotation with empty key points returns unchanged scene."""
+        annotator = sv.VertexAnnotator()
+        result = annotator.annotate(scene=scene.copy(), key_points=empty_key_points)
+
+        # Should return the original scene unchanged
+        assert np.array_equal(result, scene)
+
+
+class TestEdgeAnnotator:
+    def test_annotate_with_default_parameters(self, scene, sample_key_points):
+        """Test annotation with default parameters using COCO skeleton."""
+        annotator = sv.EdgeAnnotator()
+        result = annotator.annotate(scene=scene.copy(), key_points=sample_key_points)
+
+        # Check that the scene has been modified
+        assert_image_mostly_same(
+            original=scene, annotated=result, similarity_threshold=0.7
+        )
+
+    def test_annotate_with_custom_edges(self, scene, sample_key_points):
+        """Test annotation with custom edge definitions."""
+        edges = [(1, 2), (2, 3)]
+        annotator = sv.EdgeAnnotator(edges=edges)
+        result = annotator.annotate(scene=scene.copy(), key_points=sample_key_points)
+
+        # Check that the scene has been modified
+        assert_image_mostly_same(
+            original=scene, annotated=result, similarity_threshold=0.8
+        )
+
+    def test_annotate_empty_key_points(self, scene, empty_key_points):
+        """Test annotation with empty key points returns unchanged scene."""
+        annotator = sv.EdgeAnnotator()
+        result = annotator.annotate(scene=scene.copy(), key_points=empty_key_points)
+
+        # Should return the original scene unchanged
+        assert np.array_equal(result, scene)
+
+    def test_annotate_no_edges_found(self, scene):
+        """Test annotation when no matching skeleton is found."""
+        # Key points with more vertices than any skeleton
+        large_key_points = sv.KeyPoints(
+            xy=np.array([[[i * 10, i * 10] for i in range(100)]], dtype=np.float32),
+            confidence=np.array([[0.8] * 100], dtype=np.float32),
+            class_id=np.array([0], dtype=int),
+        )
+        annotator = sv.EdgeAnnotator()
+        result = annotator.annotate(scene=scene.copy(), key_points=large_key_points)
+
+        # Should return the original scene unchanged (no edges found)
+        assert np.array_equal(result, scene)

test/key_points/test_skeletons.py

@@ -0,0 +1,60 @@
+from supervision.key_points.skeletons import (
+    SKELETONS_BY_EDGE_COUNT,
+    SKELETONS_BY_VERTEX_COUNT,
+    Skeleton,
+)
+
+
+class TestSkeletons:
+    def test_skeleton_enum_values(self):
+        """Test skeleton enum has correct structure."""
+        for skeleton in Skeleton:
+            assert isinstance(skeleton.value, tuple)
+            assert all(
+                isinstance(edge, tuple) and len(edge) == 2 for edge in skeleton.value
+            )
+
+    def test_skeletons_by_vertex_count(self):
+        """Test SKELETONS_BY_VERTEX_COUNT dictionary population."""
+        # Test that the dictionary is populated
+        assert len(SKELETONS_BY_VERTEX_COUNT) > 0
+
+        # Test specific known skeletons
+        coco_skeleton = Skeleton.COCO.value
+        assert 17 in SKELETONS_BY_VERTEX_COUNT  # COCO has 17 keypoints
+        assert SKELETONS_BY_VERTEX_COUNT[17] == coco_skeleton
+
+    def test_skeletons_by_edge_count(self):
+        """Test SKELETONS_BY_EDGE_COUNT dictionary mapping."""
+        # Test that the dictionary is populated
+        assert len(SKELETONS_BY_EDGE_COUNT) > 0
+
+        # Reconstruct the expected mapping: for each skeleton, map the number of
+        # edges in skeleton.value to skeleton.value itself (as done in skeletons.py).
+        expected = {}
+        for skeleton in Skeleton:
+            edge_count = len(skeleton.value)
+            expected[edge_count] = skeleton.value
+
+        assert SKELETONS_BY_EDGE_COUNT == expected
+
+    def test_unique_vertices_calculation(self):
+        """Test unique vertices calculation from skeleton edges."""
+        coco_skeleton = Skeleton.COCO.value
+        unique_vertices = {vertex for edge in coco_skeleton for vertex in edge}
+        assert len(unique_vertices) == 17  # COCO has 17 keypoints
+
+    def test_skeletons_by_vertex_count_mapping_behaviour(self):
+        """Test SKELETONS_BY_VERTEX_COUNT uses last-in-wins for duplicate counts."""
+        expected_mapping = {}
+        for skeleton in Skeleton:
+            vertex_count = len({v for edge in skeleton.value for v in edge})
+            # Mimic skeletons.py: later skeletons overwrite earlier ones
+            expected_mapping[vertex_count] = skeleton.value
+
+        # The keys (vertex counts) should match
+        assert set(SKELETONS_BY_VERTEX_COUNT.keys()) == set(expected_mapping.keys())
+
+        # For each vertex count, the stored skeleton should be the last one encountered
+        for vertex_count, skeleton_value in expected_mapping.items():
+            assert SKELETONS_BY_VERTEX_COUNT[vertex_count] == skeleton_value