From 78e99354584bc95313d008feb440459960985036 Mon Sep 17 00:00:00 2001
From: Kelsey Jordahl <kjordahl@alum.mit.edu>
Date: Tue, 29 Oct 2013 19:17:29 -0400
Subject: [PATCH] TST: Refactor plot test to put repetitive code in helper
 methods

---
 tests/test_plotting.py | 41 +++++++++++++++++++++++++----------------
 1 file changed, 25 insertions(+), 16 deletions(-)

diff --git a/tests/test_plotting.py b/tests/test_plotting.py
index e3598b9..bbda9dd 100644
--- a/tests/test_plotting.py
+++ b/tests/test_plotting.py
@@ -1,4 +1,5 @@
 import os
+import shutil
 import tempfile
 import unittest
 
@@ -12,33 +13,41 @@ from geopandas import GeoSeries
 # If set to True, generate images rather than perform tests (all tests will pass!)
 GENERATE_BASELINE = False
 
-TEMPDIR = tempfile.gettempdir()
 BASELINE_DIR = os.path.join(os.path.dirname(__file__), 'baseline_images', 'test_plotting')
 
 
 class PlotTests(unittest.TestCase):
     
-    def test_poly_plot(self, tol=8):
+    def setUp(self):
+        self.tempdir = tempfile.mkdtemp()
+        return
+
+    def tearDown(self):
+        shutil.rmtree(self.tempdir)
+        return
+
+    def _compare_images(self, ax, filename, tol=8):
+        """ Helper method to do the comparisons """
+        assert isinstance(ax, Artist)
+        if GENERATE_BASELINE:
+            savefig(os.path.join(BASELINE_DIR, filename))
+        savefig(os.path.join(self.tempdir, filename))
+        err = compare_images(os.path.join(BASELINE_DIR, filename),
+                             os.path.join(self.tempdir, filename),
+                             tol, in_decorator=True)
+        if err:
+            raise ImageComparisonFailure('images not close: %(actual)s '
+                                         'vs. %(expected)s '
+                                         '(RMS %(rms).3f)' % err)
+
+    def test_poly_plot(self):
         """ Test plotting a simple series of polygons """
         filename = 'poly_plot.png'
         t1 = Polygon([(0, 0), (1, 0), (1, 1)])
         t2 = Polygon([(1, 0), (2, 0), (2, 1)])
         polys = GeoSeries([t1, t2])
         ax = polys.plot()
-        assert isinstance(ax, Artist)
-        if GENERATE_BASELINE:
-            savefig(os.path.join(BASELINE_DIR, filename))
-        savefig(os.path.join(TEMPDIR, filename))
-        err = compare_images(os.path.join(BASELINE_DIR, filename),
-                             os.path.join(TEMPDIR, filename),
-                             tol, in_decorator=True)
-        try:
-            if err:
-                raise ImageComparisonFailure('images not close: %(actual)s '
-                                             'vs. %(expected)s '
-                                             '(RMS %(rms).3f)' % err)
-        finally:
-            os.remove(os.path.join(TEMPDIR, filename))
+        self._compare_images(ax=ax, filename=filename)
 
 if __name__ == '__main__':
     unittest.main()