Add some animation tests.

anntzer · anntzer · commit 0a22a4712a41 · 2018-03-15T16:57:52.000-07:00
diff --git a/lib/matplotlib/tests/test_animation.py b/lib/matplotlib/tests/test_animation.py
@@ -1,7 +1,4 @@
-from __future__ import absolute_import, division, print_function
-
-import six
-
+from pathlib import Path
 import sys
 import tempfile
 
@@ -42,25 +39,30 @@ def finish(self):
         pass
 
 
-def test_null_movie_writer():
-    # Test running an animation with NullMovieWriter.
-
-    fig = plt.figure()
+def make_animation(**kwargs):
+    fig, ax = plt.subplots()
+    line, = ax.plot([])
 
     def init():
         pass
 
     def animate(i):
-        pass
+        line.set_data([0, 1], [0, i])
+        return line,
+
+    return animation.FuncAnimation(fig, animate, **kwargs)
+
+
+def test_null_movie_writer():
+    # Test running an animation with NullMovieWriter.
+
+    anim = make_animation(frames=5)
 
-    num_frames = 5
     filename = "unused.null"
     dpi = 50
     savefig_kwargs = dict(foo=0)
-
-    anim = animation.FuncAnimation(fig, animate, init_func=init,
-                                   frames=num_frames)
     writer = NullMovieWriter()
+
     anim.save(filename, dpi=dpi, writer=writer,
               savefig_kwargs=savefig_kwargs)
 
@@ -178,23 +180,8 @@ def animate(i):
 
 
 def test_no_length_frames():
-    fig, ax = plt.subplots()
-    line, = ax.plot([], [])
-
-    def init():
-        line.set_data([], [])
-        return line,
-
-    def animate(i):
-        x = np.linspace(0, 10, 100)
-        y = np.sin(x + i)
-        line.set_data(x, y)
-        return line,
-
-    anim = animation.FuncAnimation(fig, animate, init_func=init,
-                                   frames=iter(range(5)))
-    writer = NullMovieWriter()
-    anim.save('unused.null', writer=writer)
+    (make_animation(frames=iter(range(5)))
+     .save('unused.null', writer=NullMovieWriter()))
 
 
 def test_movie_writer_registry():
@@ -212,10 +199,30 @@ def test_movie_writer_registry():
     assert not animation.writers.is_available("ffmpeg")
     # something which is guaranteed to be available in path
     # and exits immediately
-    bin = u"true" if sys.platform != 'win32' else u"where"
+    bin = "true" if sys.platform != 'win32' else "where"
     mpl.rcParams['animation.ffmpeg_path'] = bin
     assert animation.writers._dirty
     animation.writers.list()  # resets
     assert not animation.writers._dirty
     assert animation.writers.is_available("ffmpeg")
     mpl.rcParams['animation.ffmpeg_path'] = ffmpeg_path
+
+
+@pytest.mark.parametrize("method_name", ["to_html5_video", "to_jshtml"])
+def test_embed_limit(method_name, caplog):
+    mpl.rcParams["animation.embed_limit"] = 1e-6  # Approximately 1 byte.
+    getattr(make_animation(frames=1), method_name)()
+    assert len(caplog.records) == 1
+    record, = caplog.records
+    assert (record.name == "matplotlib.animation"
+            and record.levelname == "WARNING")
+
+
+@pytest.mark.parametrize(
+    "method_name",
+    ["to_html5_video",
+     pytest.mark.xfail("to_jshtml")])  # Needs to be fixed.
+def test_cleanup_temporaries(method_name, tmpdir):
+    with tmpdir.as_cwd():
+        getattr(make_animation(frames=1), method_name)()
+        assert list(Path(tmpdir).iterdir()) == []
