Add test for category colorbar calls

trollimage/colormap.py

@@ -1078,14 +1078,15 @@ def _read_colormap_data_from_np(path):
 def colorbar(height, length, colormap, category=False):
     """Return the channels of a colorbar."""
     cbar = np.tile(np.arange(length) * 1.0 / (length - 1), (height, 1))
+    cmin = colormap.values.min()
+    cmax = colormap.values.max()
+    crange = (cmax - cmin)
     if category:
         # add an extra buffer around colormap limits to show full category
-        cbar = (cbar * (colormap.values.max() - colormap.values.min() + 1)
-                + colormap.values.min())
-        cbar = np.round(cbar - 0.5)
+        cbar = cbar * (crange + 1) + (cmin - 0.5)
+        cbar = np.round(cbar)
     else:
-        cbar = (cbar * (colormap.values.max() - colormap.values.min())
-                + colormap.values.min())
+        cbar = (cbar * crange) + colormap.values.min()
 
     return colormap.colorize(cbar)
 

trollimage/tests/test_colormap.py

@@ -102,16 +102,23 @@ def test_add(self):
         np.testing.assert_allclose(cm3.colors, cm_.colors)
         np.testing.assert_allclose(cm3.values, cm_.values)
 
-    def test_colorbar(self):
+    @pytest.mark.parametrize("category", [False, True])
+    def test_colorbar(self, category):
         """Test colorbar."""
         cm_ = colormap.Colormap((1, (1.0, 1.0, 0.0)),
                                 (2, (0.0, 1.0, 1.0)),
-                                (3, (1.0, 1.0, 1.0)),
-                                (4, (0.0, 0.0, 0.0)))
-
-        channels = colormap.colorbar(1, 4, cm_)
-        for i in range(3):
-            np.testing.assert_allclose(channels[i].ravel(), cm_.colors[:, i], atol=0.001)
+                                (3, (1.0, 0.0, 1.0)),
+                                (4, (1.0, 1.0, 1.0)),
+                                (5, (0.0, 0.0, 0.0)))
+
+        channels = colormap.colorbar(1, 9, cm_, category=category)
+        channels = np.array(channels)[:, 0, :]
+        # number of colors and size of colorbar specifically chosen to have exact
+        # colors of the colormap show in the colorbar
+        assert np.unique(channels, axis=1).shape[1] == (5 if category else 9)
+        for exp_color in cm_.colors:
+            # check that the colormaps colors are actually showing up
+            assert np.isclose(channels, exp_color[:, None], atol=0.01).all(axis=0).any()
 
     def test_palettebar(self):
         """Test colorbar."""