Test set_array for copy and wrong dtype

aitikgupta · aitikgupta · commit a071304f8bd7 · 2021-04-29T07:23:30.000+05:30
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -1,5 +1,6 @@
 import io
 from types import SimpleNamespace
+from typing import Type
 
 import numpy as np
 from numpy.testing import assert_array_equal, assert_array_almost_equal
@@ -677,6 +678,22 @@ def test_collection_set_verts_array():
         assert np.array_equal(ap._codes, atp._codes)
 
 
+def test_collection_set_array():
+    vals = [*range(10)]
+
+    # Test set_array with list
+    c = Collection()
+    c.set_array(vals)
+
+    # Test set_array with wrong dtype
+    with pytest.raises(TypeError, match="^Image data of dtype"):
+        c.set_array("wrong_input")
+
+    # Test if array kwarg is copied
+    vals[5] = 45
+    assert np.not_equal(vals, c.get_array()).any()
+
+
 def test_blended_collection_autolim():
     a = [1, 2, 4]
     height = .2
