diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 6c3b8a2..1e80c0e 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -15,7 +15,7 @@ jobs:
         - macos: py39-test
         - windows: py39-test
         - linux: py310-test
-        - macos: py311-test
+        - macos: py311-test-viewer
         - windows: py312-test
         - linux: py313-test-devdeps
       coverage: 'codecov'
diff --git a/astrodendro/tests/test_viewer.py b/astrodendro/tests/test_viewer.py
new file mode 100644
index 0000000..8d68487
--- /dev/null
+++ b/astrodendro/tests/test_viewer.py
@@ -0,0 +1,41 @@
+import pytest
+import numpy as np
+
+import matplotlib.pyplot as plt
+from ..dendrogram import Dendrogram
+from matplotlib.backend_bases import MouseEvent
+
+
+DATA = np.array([[1, 3, 4, 4, 1, 4],
+                 [1, 2, 3, 2, 1, 3],
+                 [2, 1, 1, 3, 1, 2],
+                 [1, 1, 1, 1, 1, 1],
+                 [2, 3, 2, 1, 1, 2],
+                 [2, 3, 5, 3, 1, 1]])
+
+
+def test_viewer(capsys):
+
+    original_backend = plt.get_backend()
+
+    try:
+        plt.switch_backend('qtagg')
+    except ImportError:
+        pytest.skip("This test requires Qt to be installed")
+
+    d = Dendrogram.compute(DATA)
+    viewer = d.viewer()
+
+    plt.show(block=False)
+
+    cb = viewer.fig.canvas.callbacks
+
+    cb.process('button_press_event', MouseEvent('button_press_event', viewer.fig.canvas, 660, 520, 1))
+    cb.process('button_press_event', MouseEvent('button_press_event', viewer.fig.canvas, 890, 800, 1))
+    cb.process('button_press_event', MouseEvent('button_press_event', viewer.fig.canvas, 700, 700, 1))
+
+    plt.switch_backend(original_backend)
+
+    captured = capsys.readouterr()
+    assert captured.out == ""
+    assert captured.err == ""
diff --git a/astrodendro/viewer.py b/astrodendro/viewer.py
index 49822a3..b835bb4 100644
--- a/astrodendro/viewer.py
+++ b/astrodendro/viewer.py
@@ -277,7 +277,7 @@ def _update_lines(self, selection_id):
 
         # Remove previously selected collection
         if selection_id in self.selected_lines:
-            self.ax_dendrogram.collections.remove(self.selected_lines[selection_id])
+            self.selected_lines[selection_id].remove()
             del self.selected_lines[selection_id]
 
         if structure is None:
@@ -308,15 +308,13 @@ def _update_lines(self, selection_id):
         self.ax_dendrogram.add_collection(self.selected_lines[selection_id])
 
     def remove_contour(self, selection_id):
-
         if selection_id in self.selected_contour:
-            for collection in self.selected_contour[selection_id].collections:
-                self.ax_image.collections.remove(collection)
+            self.selected_contour[selection_id].remove()
             del self.selected_contour[selection_id]
 
     def remove_all_contours(self):
         """ Remove all selected contours. """
-        for key in self.selected_contour.keys():
+        for key in list(self.selected_contour):
             self.remove_contour(key)
 
     def update_contours(self):
diff --git a/tox.ini b/tox.ini
index c344359..a8b89f7 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,6 +1,6 @@
 [tox]
 envlist =
-    py{38,39,310,311}-test{,-alldeps,-devdeps}{,-cov}
+    py{38,39,310,311}-test{,-alldeps,-devdeps,-viewer}{,-cov}
     build_docs
     linkcheck
     codestyle
@@ -10,6 +10,8 @@ requires =
 isolated_build = true
 
 [testenv]
+passenv =
+    DISPLAY
 setenv =
     MPLBACKEND=agg
     devdeps: PIP_EXTRA_INDEX_URL = https://pypi.anaconda.org/astropy/simple https://pypi.anaconda.org/scientific-python-nightly-wheels/
@@ -30,6 +32,7 @@ deps =
     oldestdeps: matplotlib==3.3.*
     oldestdeps: numpy==1.20.*
     oldestdeps: pillow==8.0.*
+    viewer: PyQt6
 
 extras =
     test