Add test for quiverkey with multiple angles

dstansby · dstansby · commit 2c2ff2d33b21 · 2017-01-01T15:50:06.000Z
diff --git a/lib/matplotlib/tests/test_quiver.py b/lib/matplotlib/tests/test_quiver.py
@@ -161,6 +161,20 @@ def test_bad_masked_sizes():
     ax.barbs(x, y, u, v)
 
 
+def test_quiverkey_angles():
+    # Check that only a single arrow is plotted for a quiverkey when an array
+    # of angles is given to the original quiver plot
+    fig, ax = plt.subplots()
+
+    X, Y = np.meshgrid(np.arange(2), np.arange(2))
+    U = V = angles = np.ones_like(X)
+
+    q = ax.quiver(X, Y, U, V, angles=angles)
+    qk = ax.quiverkey(q, 1, 1, 2, 'Label')
+    # The arrows are only created when the key is drawn
+    fig.canvas.draw()
+    assert len(qk.verts) == 1
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()
