diff --git a/tests/_images/Labels_can_color_labels.png b/tests/_images/Labels_can_color_labels.png
index 538ef4a2..862b5d88 100644
Binary files a/tests/_images/Labels_can_color_labels.png and b/tests/_images/Labels_can_color_labels.png differ
diff --git a/tests/_images/Labels_label_categorical_color.png b/tests/_images/Labels_label_categorical_color.png
index 7d59893d..2a88dcfa 100644
Binary files a/tests/_images/Labels_label_categorical_color.png and b/tests/_images/Labels_label_categorical_color.png differ
diff --git a/tests/pl/test_render_labels.py b/tests/pl/test_render_labels.py
index cab4e6cf..28955398 100644
--- a/tests/pl/test_render_labels.py
+++ b/tests/pl/test_render_labels.py
@@ -111,12 +111,13 @@ def test_plot_label_categorical_color(self, sdata_blobs: SpatialData, label: str
         self._make_tablemodel_with_categorical_labels(sdata_blobs, label)
 
     def _make_tablemodel_with_categorical_labels(self, sdata_blobs, label):
-        n_obs = max(get_element_instances(sdata_blobs[label]))
+        instances = get_element_instances(sdata_blobs[label])
+        n_obs = len(instances)
         adata = AnnData(
             RNG.normal(size=(n_obs, 10)),
             obs=pd.DataFrame(RNG.normal(size=(n_obs, 3)), columns=["a", "b", "c"]),
         )
-        adata.obs["instance_id"] = np.arange(adata.n_obs)
+        adata.obs["instance_id"] = instances.values
         adata.obs["category"] = RNG.choice(["a", "b", "c"], size=adata.n_obs)
         adata.obs["category"][:3] = ["a", "b", "c"]
         adata.obs["instance_id"] = list(range(adata.n_obs))