From a0a47b34a3a158d8172e9ef2c050b59287e696ee Mon Sep 17 00:00:00 2001
From: Genevieve Buckley <30920819+GenevieveBuckley@users.noreply.github.com>
Date: Tue, 17 Oct 2023 20:07:38 +1100
Subject: [PATCH] Simple GUI test for CI (#223)

Enable GUI tests in CLI and add a simple GUI test.
---
 .github/workflows/test.yaml | 16 ++++++++++---
 micro_sam/visualization.py  |  1 +
 pyproject.toml              |  6 ++++-
 test/test_gui.py            | 48 +++++++++++++++++++++++++++++++++++++
 4 files changed, 67 insertions(+), 4 deletions(-)
 create mode 100644 test/test_gui.py

diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index fcc4263c..da1de6ff 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -15,7 +15,7 @@ jobs:
   test:
     name: ${{ matrix.os }} ${{ matrix.python-version }}
     runs-on: ${{ matrix.os }}
-
+    timeout-minutes: 60
     strategy:
       fail-fast: false
       matrix:
@@ -31,19 +31,29 @@ jobs:
         with:
           environment-file: environment_cpu.yaml
 
+      # Setup Qt libraries for GUI testing
+      - uses: tlambert03/setup-qt-libs@v1
+
+      - name: Install napari and pyqt for GUI tests
+        shell: bash -l {0}
+        run: pip install "napari[all]"
+
       - name: Install pytest
         shell: bash -l {0}
         run: |
           python -m pip install pytest
           python -m pip install pytest-cov
+          python -m pip install pytest-qt
 
       - name: Install package
         shell: bash -l {0}
         run: pip install --no-deps -e .
 
       - name: Run tests
-        shell: bash -l {0}
-        run: pytest
+        uses: aganders3/headless-gui@v1
+        with:
+          shell: bash -l {0}
+          run: pytest
 
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v3
diff --git a/micro_sam/visualization.py b/micro_sam/visualization.py
index 0426e193..c931985f 100644
--- a/micro_sam/visualization.py
+++ b/micro_sam/visualization.py
@@ -1,6 +1,7 @@
 """
 Functionality for visualizing image embeddings.
 """
+from typing import Tuple
 
 import numpy as np
 
diff --git a/pyproject.toml b/pyproject.toml
index d636f5bb..1ce35bd9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,11 @@ requires = ["setuptools>=42.0.0", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [tool.pytest.ini_options]
-addopts = "-v --cov=micro_sam --cov-report xml:coverage.xml --cov-report term"
+addopts = "-v --durations=10 --cov=micro_sam --cov-report xml:coverage.xml"
+markers = [
+    "slow: marks tests as slow (deselect with '-m \"not slow\"')",
+    "gui: marks GUI tests (deselect with '-m \"not gui\"')",
+]
 
 [tool.black]
 line-length = 79
diff --git a/test/test_gui.py b/test/test_gui.py
new file mode 100644
index 00000000..e82f5610
--- /dev/null
+++ b/test/test_gui.py
@@ -0,0 +1,48 @@
+import numpy as np
+import pytest
+
+from micro_sam.sam_annotator import annotator_2d
+from micro_sam.sam_annotator.annotator_2d import _initialize_viewer
+
+
+def _check_layer_initialization(viewer):
+    """Utility function to check the initial layer setup is correct."""
+    assert len(viewer.layers) == 6
+    expected_layer_names = ['raw', 'auto_segmentation', 'committed_objects', 'current_object', 'point_prompts', 'prompts']
+    for layername in expected_layer_names:
+        assert layername in viewer.layers
+    # Check layers are empty before beginning tests
+    np.testing.assert_equal(viewer.layers["auto_segmentation"].data, 0)
+    np.testing.assert_equal(viewer.layers["current_object"].data, 0)
+    np.testing.assert_equal(viewer.layers["committed_objects"].data, 0)
+    np.testing.assert_equal(viewer.layers["point_prompts"].data, 0)
+    assert viewer.layers["prompts"].data == []  # shape data is list, not numpy array
+
+
+@pytest.mark.gui
+def test_annotator_2d(make_napari_viewer_proxy, tmp_path):
+    """Integration test for annotator_2d widget with automatic mask generation.
+
+    * Creates 2D image embedding
+    * Opens annotator_2d widget in napari
+    """
+    model_type = "vit_b"
+    embedding_path = tmp_path / "test-embedding.zarr"
+    # example data - a basic checkerboard pattern
+    image = np.zeros((16,16))
+    image[:8,:8] = 1
+    image[8:,8:] = 1
+
+    viewer = make_napari_viewer_proxy()
+    viewer = _initialize_viewer(image, None, None, None)  # TODO: fix hacky workaround
+    # test generating image embedding, then adding micro-sam dock widgets to the GUI
+    viewer = annotator_2d(
+        image,
+        embedding_path,
+        show_embeddings=False,
+        model_type=model_type,
+        v=viewer,
+        return_viewer=True
+    )
+    _check_layer_initialization(viewer)
+    viewer.close()  # must close the viewer at the end of tests