Use MagicMocks for headless display tests

Replace handcrafted FakeTk/Label/PhotoImage classes with unittest.mock MagicMocks in the headless_display_env fixture. The fixture now returns a SimpleNamespace with mock constructors/instances (tk, label, photo) and sets display module attributes with monkeypatch (raising=False). Update tests to use the env mocks, assert on mock calls (title, pack, configure, update, destroy), mock ImageDraw.Draw with a MagicMock, and make colormap sampling deterministic for assertions. This simplifies test setup and enables precise call-based assertions instead of inspecting custom fake object state.

Author: C-Achard

tests/conftest.py

@@ -2,7 +2,9 @@
 
 import copy
 from collections.abc import Callable
+from types import SimpleNamespace
 from typing import Any
+from unittest.mock import MagicMock
 
 import numpy as np
 import pytest
@@ -15,20 +17,50 @@
 # --------------------------------------------------------------------------------------
 @pytest.fixture
 def headless_display_env(monkeypatch):
-    """Patch dlclive.display so tkinter is replaced with fake, non-GUI-safe objects."""
-    from test_display import FakeLabel, FakePhotoImage, FakeTk
-
+    """
+    Patch dlclive.display so tkinter + ImageTk are replaced by MagicMocks.
+
+    Returns an object with:
+      - mod: the imported dlclive.display module
+      - tk_ctor: MagicMock constructor for Tk
+      - tk: MagicMock instance for the window
+      - label_ctor: MagicMock constructor for Label
+      - label: MagicMock instance for the label widget
+      - photo_ctor: MagicMock function for ImageTk.PhotoImage
+      - photo: MagicMock instance representing created image
+    """
     import dlclive.display as display_mod
 
-    monkeypatch.setattr(display_mod, "_TKINTER_AVAILABLE", True)
-    monkeypatch.setattr(display_mod, "Tk", FakeTk)
-    monkeypatch.setattr(display_mod, "Label", FakeLabel)
+    # Ensure display path is enabled
+    monkeypatch.setattr(display_mod, "_TKINTER_AVAILABLE", True, raising=False)
 
-    class FakeImageTkModule:
-        PhotoImage = FakePhotoImage
+    # Tk / Label mocks
+    tk = MagicMock(name="TkInstance")
+    tk_ctor = MagicMock(name="Tk", return_value=tk)
 
-    monkeypatch.setattr(display_mod, "ImageTk", FakeImageTkModule)
-    return display_mod
+    label = MagicMock(name="LabelInstance")
+    label_ctor = MagicMock(name="Label", return_value=label)
+
+    # ImageTk.PhotoImage mock
+    photo = MagicMock(name="PhotoImageInstance")
+    photo_ctor = MagicMock(name="PhotoImage", return_value=photo)
+
+    class FakeImageTkModule:
+        PhotoImage = photo_ctor
+
+    monkeypatch.setattr(display_mod, "Tk", tk_ctor, raising=False)
+    monkeypatch.setattr(display_mod, "Label", label_ctor, raising=False)
+    monkeypatch.setattr(display_mod, "ImageTk", FakeImageTkModule, raising=False)
+
+    return SimpleNamespace(
+        mod=display_mod,
+        tk_ctor=tk_ctor,
+        tk=tk,
+        label_ctor=label_ctor,
+        label=label,
+        photo_ctor=photo_ctor,
+        photo=photo,
+    )
 
 
 # --------------------------------------------------------------------------------------

tests/test_display.py

@@ -1,42 +1,9 @@
+from unittest.mock import ANY, MagicMock
+
 import numpy as np
 import pytest
 
 
-class FakeTk:
-    def __init__(self):
-        self.titles = []
-        self.updated = 0
-        self.destroyed = False
-
-    def title(self, text):
-        self.titles.append(text)
-
-    def update(self):
-        self.updated += 1
-
-    def destroy(self):
-        self.destroyed = True
-
-
-class FakeLabel:
-    def __init__(self, window):
-        self.window = window
-        self.packed = False
-        self.configured = {}
-
-    def pack(self):
-        self.packed = True
-
-    def configure(self, **kwargs):
-        self.configured.update(kwargs)
-
-
-class FakePhotoImage:
-    def __init__(self, image=None, master=None):
-        self.image = image
-        self.master = master
-
-
 def test_display_init_raises_when_tk_unavailable(monkeypatch):
     import dlclive.display as display_mod
 
@@ -47,52 +14,66 @@ def test_display_init_raises_when_tk_unavailable(monkeypatch):
 
 
 def test_display_frame_creates_window_and_updates(headless_display_env):
-    display_mod = headless_display_env
+    env = headless_display_env
+    display_mod = env.mod
     disp = display_mod.Display(radius=3, pcutoff=0.5)
 
     frame = np.zeros((100, 120, 3), dtype=np.uint8)
     pose = np.array([[[10, 10, 0.9], [50, 50, 0.2]]])  # 1 animal, 2 bodyparts
 
     disp.display_frame(frame, pose)
 
-    assert disp.window is not None
-    assert disp.lab is not None
-    assert disp.lab.packed is True
-    assert disp.window.updated == 1
-    assert "image" in disp.lab.configured  # configured with PhotoImage
+    # Window created and initialized
+    env.tk_ctor.assert_called_once_with()
+    env.tk.title.assert_called_once_with("DLC Live")
+
+    # Label created and packed
+    env.label_ctor.assert_called_once_with(env.tk)
+    env.label.pack.assert_called_once()
+
+    # PhotoImage created with correct master + image passed
+    env.photo_ctor.assert_called_once_with(image=ANY, master=env.tk)
+
+    # Image configured on label and window updated
+    env.label.configure.assert_called_once_with(image=env.photo)
+    env.tk.update.assert_called_once_with()
 
 
 def test_display_draws_only_points_above_cutoff(headless_display_env, monkeypatch):
-    display_mod = headless_display_env
+    env = headless_display_env
+    display_mod = env.mod
     disp = display_mod.Display(radius=3, pcutoff=0.5)
 
+    # Patch colormap so color sampling is deterministic and always long enough
+    class FakeCC:
+        bmy = [(1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0)]
+
+    monkeypatch.setattr(display_mod, "cc", FakeCC)
+
     frame = np.zeros((100, 100, 3), dtype=np.uint8)
     pose = np.array(
         [
             [
                 [10, 10, 0.9],  # draw
                 [20, 20, 0.49],  # don't draw
-                [30, 30, 0.5001],  # draw (>=)
+                [30, 30, 0.5001],  # draw (> pcutoff)
             ]
         ],
         dtype=float,
     )
 
-    ellipses = []
-
-    class DrawRecorder:
-        def ellipse(self, coords, fill=None, outline=None):
-            ellipses.append((coords, fill, outline))
-
-    monkeypatch.setattr(display_mod.ImageDraw, "Draw", lambda img: DrawRecorder())
+    draw = MagicMock(name="DrawInstance")
+    monkeypatch.setattr(display_mod.ImageDraw, "Draw", MagicMock(return_value=draw))
 
     disp.display_frame(frame, pose)
 
-    assert len(ellipses) == 2
+    # Two points above cutoff => two ellipse calls
+    assert draw.ellipse.call_count == 2
 
 
 def test_destroy_calls_window_destroy(headless_display_env):
-    display_mod = headless_display_env
+    env = headless_display_env
+    display_mod = env.mod
     disp = display_mod.Display()
 
     frame = np.zeros((10, 10, 3), dtype=np.uint8)
@@ -101,13 +82,13 @@ def test_destroy_calls_window_destroy(headless_display_env):
     disp.display_frame(frame, pose)
     disp.destroy()
 
-    assert disp.window.destroyed is True
+    env.tk.destroy.assert_called_once_with()
 
 
 def test_set_display_color_sampling_safe(headless_display_env, monkeypatch):
-    display_mod = headless_display_env
+    env = headless_display_env
+    display_mod = env.mod
 
-    # Provide a fixed colormap list
     class FakeCC:
         bmy = [(1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), (0, 1, 1), (1, 0, 1)]
 
@@ -118,3 +99,9 @@ class FakeCC:
 
     assert disp.colors is not None
     assert len(disp.colors) >= 3
+
+    # Also verify window setup calls happened
+    env.tk_ctor.assert_called_once_with()
+    env.tk.title.assert_called_once_with("DLC Live")
+    env.label_ctor.assert_called_once_with(env.tk)
+    env.label.pack.assert_called_once()