fix plotting tests: the effective colorbar is not anymore necessarily at the first collection position

Author: emanuel-schmid

climada/util/test/test_plot.py

@@ -100,8 +100,9 @@ def test_geo_scatter_from_array(self):
                                shapes=True, axes=None, proj=projection,
                                figsize=(9, 13), cmap=cmap)
         self.assertEqual(var_name, ax.get_title())
-        self.assertAlmostEqual(np.max(values), ax.collections[0].colorbar.vmax)
-        self.assertAlmostEqual(np.min(values), ax.collections[0].colorbar.vmin)
+        colorbar = next(x.colorbar for x in ax.collections if x.colorbar)
+        self.assertAlmostEqual(np.max(values), colorbar.vmax)
+        self.assertAlmostEqual(np.min(values), colorbar.vmin)
         self.assertEqual(cmap, ax.collections[0].cmap.name)
         plt.close()
 
@@ -117,8 +118,9 @@ def test_geo_bin_from_array(self):
                                            shapes=True, axes=None, proj=projection,
                                            figsize=(9, 13), cmap=cmap)
         self.assertEqual(var_name, ax.get_title())
-        self.assertAlmostEqual(np.max(values), ax.collections[0].colorbar.vmax)
-        self.assertAlmostEqual(np.min(values), ax.collections[0].colorbar.vmin)
+        colorbar = next(x.colorbar for x in ax.collections if x.colorbar)
+        self.assertAlmostEqual(np.max(values), colorbar.vmax)
+        self.assertAlmostEqual(np.min(values), colorbar.vmin)
         self.assertEqual(cmap, ax.collections[0].cmap.name)
         plt.close()
 
@@ -132,17 +134,19 @@ def test_geo_im_from_array(self):
         ax = u_plot.geo_im_from_array(values, coord, var_name, title,
                       proj=projection, smooth=True, axes=None, figsize=(9, 13), cmap=cmap)
         self.assertEqual(var_name, ax.get_title())
-        self.assertAlmostEqual(np.max(values), ax.collections[0].colorbar.vmax)
-        self.assertAlmostEqual(np.min(values), ax.collections[0].colorbar.vmin)
+        colorbar = next(x.colorbar for x in ax.collections if x.colorbar)
+        self.assertAlmostEqual(np.max(values), colorbar.vmax)
+        self.assertAlmostEqual(np.min(values), colorbar.vmin)
         self.assertEqual(cmap, ax.collections[0].cmap.name)
         plt.close()
 
         projection = ccrs.AzimuthalEquidistant()
         ax = u_plot.geo_im_from_array(values, coord, var_name, title,
                       proj=projection, smooth=True, axes=None, figsize=(9, 13), cmap=cmap)
         self.assertEqual(var_name, ax.get_title())
-        self.assertAlmostEqual(np.max(values), ax.collections[0].colorbar.vmax)
-        self.assertAlmostEqual(np.min(values), ax.collections[0].colorbar.vmin)
+        colorbar = next(x.colorbar for x in ax.collections if x.colorbar)
+        self.assertAlmostEqual(np.max(values), colorbar.vmax)
+        self.assertAlmostEqual(np.min(values), colorbar.vmin)
         self.assertEqual(cmap, ax.collections[0].cmap.name)
         plt.close()