Add tests for ebsd.Map.find_grains

tests/test_ebsd.py

@@ -1,6 +1,7 @@
 import pytest
 from pytest import approx
-from unittest.mock import Mock
+from unittest.mock import Mock, MagicMock
+from functools import partial
 
 import numpy as np
 import defdap.ebsd as ebsd
@@ -38,6 +39,18 @@ def good_phase_array(good_map_with_quats):
     return good_map_with_quats.data.phase
 
 
+@pytest.fixture(scope="module")
+def good_grain_boundaries(good_map):
+    expected = np.load(
+        f'{EXPECTED_RESULTS_DIR}/ebsd_grain_boundaries_5deg.npz'
+    )
+    return ebsd.BoundarySet(
+        good_map,
+        [tuple(row) for row in expected['points_x']],
+        [tuple(row) for row in expected['points_y']]
+    )
+
+
 @pytest.fixture(scope="module")
 def good_symmetries():
     over_root2 = np.sqrt(2) / 2
@@ -116,15 +129,105 @@ def test_calc(mock_map, bound_def):
         )
 
         # load expected
-        expected = np.loadtxt(
-            "{:}boundaries_{:}deg.txt".format(EXPECTED_RESULTS_DIR, bound_def),
-            dtype=int
+        expected = np.load(
+            f'{EXPECTED_RESULTS_DIR}/ebsd_grain_boundaries_{bound_def}deg.npz'
         )
+        expected_x = set([tuple(row) for row in expected['points_x']])
+        expected_y = set([tuple(row) for row in expected['points_y']])
+
+        assert result.points_x == expected_x
+        assert result.points_y == expected_y
+
+
+class TestMapFindGrains:
+    # Depends on self.data.grain_boundaries.image_*, self.data.phase,
+    # self.flood_fill, self.num_phases, self.phases
+    # Affects self.boundaries
+
+    @staticmethod
+    @pytest.fixture
+    def mock_map(good_grain_boundaries, good_phase_array):
+        # create stub object
+        mock_map = Mock(spec=ebsd.Map)
+        mock_datastore = MagicMock(spec=Datastore)
+        mock_datastore.phase = good_phase_array
+        mock_datastore.grain_boundaries = good_grain_boundaries
+        mock_datastore.generate_id = Mock(return_value=1)
+        # mock_datastore.__iter__.return_value = []
+        mock_map.data = mock_datastore
+        mock_map.shape = good_phase_array.shape
+        mock_map.flood_fill = partial(ebsd.Map.flood_fill, mock_map)
+        mock_map.num_phases = 1
+        mock_map.phases = [Mock(crystal.Phase)]
+
+        return mock_map
+
+    @staticmethod
+    def test_return_type(mock_map):
+        # run test and collect result
+        result = ebsd.Map.find_grains(mock_map, min_grain_size=10)
 
-        assert np.allclose(result.image, expected)
+        assert isinstance(result, np.ndarray)
+        assert result.shape == mock_map.shape
+        assert result.dtype == np.int64
+
+    @staticmethod
+    @pytest.mark.parametrize('min_grain_size', [0, 10, 100])
+    def test_calc(mock_map, min_grain_size):
+        # run test and collect result
+        result = ebsd.Map.find_grains(mock_map, min_grain_size=min_grain_size)
+
+        # load expected
+        expected = np.load(
+            f'{EXPECTED_RESULTS_DIR}/ebsd_grains_5deg_{min_grain_size}.npz'
+        )['grains']
+
+        assert np.alltrue(result == expected)
+
+    @staticmethod
+    def test_add_derivative(mock_map):
+        mock_add_derivative = Mock()
+        mock_map.data.add_derivative = mock_add_derivative
+        # run test and collect result
+        result = ebsd.Map.find_grains(mock_map, min_grain_size=10)
+
+        mock_add_derivative.assert_called_once()
+
+    @staticmethod
+    def test_grain_list_type(mock_map):
+        ebsd.Map.find_grains(mock_map, min_grain_size=10)
+        result = mock_map._grains
+
+        assert isinstance(result, list)
+        for g in result:
+            assert isinstance(g, ebsd.Grain)
+
+    @staticmethod
+    @pytest.mark.parametrize('min_grain_size, expected_len', [
+        (0, 141), (10, 109), (100, 76)
+    ])
+    def test_grain_list_size(mock_map, min_grain_size, expected_len):
+        ebsd.Map.find_grains(mock_map, min_grain_size=min_grain_size)
+        result = mock_map._grains
+
+        assert len(result) == expected_len
+
+    @staticmethod
+    @pytest.mark.parametrize('min_grain_size', [0, 10, 100])
+    def test_grain_points(mock_map, min_grain_size):
+        ebsd.Map.find_grains(mock_map, min_grain_size=min_grain_size)
+        result = mock_map._grains
+
+        # load expected
+        expected_grains = np.load(
+            f'{EXPECTED_RESULTS_DIR}/ebsd_grains_5deg_{min_grain_size}.npz'
+        )['grains']
 
+        for i in range(expected_grains.max()):
 
+            expected_point = zip(*np.nonzero(expected_grains == i+1)[::-1])
 
+            assert set(result[i].data.point) == set(expected_point)
 
 ''' Functions left to test
 Map:
@@ -146,7 +249,6 @@ def test_calc(mock_map, bound_def):
 findPhaseBoundaries
 plot_phase_boundary_map
 plot_boundary_map
-findGrains
 plot_grain_map
 floodFill
 calc_grain_av_oris

tests/test_quat.py

@@ -9,8 +9,7 @@
 # Initialisation tests
 
 
-@pytest.mark.parametrize('inputLength',
-                         [2, 5, 6])
+@pytest.mark.parametrize('inputLength', [2, 5, 6])
 def testInitDimension(inputLength):
     """Quat initialisation should raise a DimensionError if not length
     1, 3, 4"""