From 64249a66d15a60d14630c8ee28969e5bb05014a8 Mon Sep 17 00:00:00 2001
From: yrkim98 <yrkim98@uw.edu>
Date: Wed, 18 Oct 2023 11:48:01 -0700
Subject: [PATCH] change tests to use service for curation_main

---
 .../_tests/curation/test_curation_main.py     | 44 ++++++++++++++++++-
 .../curation/main_view.py                     |  6 ---
 2 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/src/allencell_ml_segmenter/_tests/curation/test_curation_main.py b/src/allencell_ml_segmenter/_tests/curation/test_curation_main.py
index c61b8283..2db3a03a 100644
--- a/src/allencell_ml_segmenter/_tests/curation/test_curation_main.py
+++ b/src/allencell_ml_segmenter/_tests/curation/test_curation_main.py
@@ -6,7 +6,6 @@
 from allencell_ml_segmenter.curation.curation_model import CurationModel
 from allencell_ml_segmenter.curation.curation_service import CurationService
 
-import napari
 from unittest.mock import Mock
 from pytestqt.qtbot import QtBot
 from pathlib import Path
@@ -19,6 +18,21 @@ def curation_main_view(qtbot: QtBot) -> CurationMainView:
     curation_service: Mock = Mock(spec=CurationService)
     return CurationMainView(curation_model, experiments_model, curation_service)
 
+def test_curation_setup(curation_main_view):
+    # Arrange
+    curation_main_view._curation_service.get_raw_images_list = Mock(return_value=[Path("path_raw")])
+    curation_main_view._curation_service.get_seg1_images_list = Mock(return_value=[Path("path_seg1")])
+    curation_main_view.init_progress_bar = Mock()
+
+    # Act
+    curation_main_view.curation_setup()
+
+    # Assert
+    curation_main_view.init_progress_bar.assert_called_once()
+    curation_main_view._curation_service.remove_all_images_from_viewer_layers.assert_called_once()
+    curation_main_view._curation_service.add_image_to_viewer.called_once_with(["path_raw"], f"[raw] path_raw")
+    curation_main_view._curation_service.add_image_to_viewer.called_once_with(["path_seg1"], f"[raw] path_seg1")
+
 
 def test_curation_main_view_init(curation_main_view: CurationMainView) -> None:
     # Arrange, Act, Assert
@@ -32,6 +46,34 @@ def test_init_progress_bar(curation_main_view: CurationMainView) -> None:
     # Assert
     assert curation_main_view.progress_bar.value() == 1
 
+def test_next_image(curation_main_view: CurationMainView) -> None:
+    # Arrange
+    curation_main_view._update_curation_record = Mock()
+    curation_main_view.raw_images = [None, Path("path_raw")]
+    curation_main_view.seg1_images = [None, Path("path_seg1")]
+    assert curation_main_view.curation_index == 0
+
+    # Act
+    curation_main_view.next_image()
+
+    # Assert
+    curation_main_view._update_curation_record.assert_called_once()
+    assert curation_main_view.curation_index == 1
+    curation_main_view._curation_service.remove_all_images_from_viewer_layers.assert_called_once()
+    assert curation_main_view._curation_service.add_image_to_viewer.call_count == 2
+
+def test_next_end_of_curation(curation_main_view: CurationMainView) -> None:
+    # Arrange
+    curation_main_view._update_curation_record = Mock()
+    curation_main_view.raw_images = [None, Path("path_raw")]
+    curation_main_view.seg1_images = [None, Path("path_seg1")]
+    curation_main_view.curation_index = 1
+
+    # Act
+    curation_main_view.next_image()
+
+    # Assert
+    assert curation_main_view._curation_service.write_curation_record.assert_called_once()
 
 def test_increment_progress_bar(curation_main_view: CurationMainView) -> None:
     # Arrange
diff --git a/src/allencell_ml_segmenter/curation/main_view.py b/src/allencell_ml_segmenter/curation/main_view.py
index b629d2b5..9fc91bbb 100644
--- a/src/allencell_ml_segmenter/curation/main_view.py
+++ b/src/allencell_ml_segmenter/curation/main_view.py
@@ -1,7 +1,5 @@
-import csv
 from pathlib import Path
 from typing import List
-import napari
 from qtpy.QtCore import Qt
 from qtpy.QtWidgets import (
     QVBoxLayout,
@@ -13,8 +11,6 @@
     QProgressBar,
     QRadioButton,
 )
-from aicsimageio import AICSImage
-
 from allencell_ml_segmenter._style import Style
 from allencell_ml_segmenter.core.view import View
 from allencell_ml_segmenter.curation.curation_data_class import CurationRecord
@@ -24,8 +20,6 @@
 from allencell_ml_segmenter.widgets.label_with_hint_widget import LabelWithHint
 
 from napari.utils.notifications import show_info
-from napari.layers.shapes.shapes import Shapes
-import shutil
 
 
 class CurationMainView(View):