test_template_to_temp_file tries to clean up after itself.

1 files changed, 13 insertions, 5 deletions

diff --git a/wkhtmltopdf/tests.py b/wkhtmltopdf/tests.py
index 5eb733f..0249cb5 100644
--- a/wkhtmltopdf/tests.py
+++ b/wkhtmltopdf/tests.py
@@ -1,13 +1,21 @@
+from __future__ import absolute_import
+
+import os
+
 from django.test import TestCase
 
-from utils import template_to_temp_file
+from .utils import template_to_temp_file
+
 
 class TestUtils(TestCase):
     def test_template_to_temp_file(self):
         """Should render a template to a temporary file."""
         title = 'A test template.'
         temp_file = template_to_temp_file('sample.html', {'title': title})
-        with open(temp_file, 'r') as f:
-            saved_content = f.read()
-        self.assertTrue(title in saved_content)
-
+        try:
+            with open(temp_file, 'r') as f:
+                saved_content = f.read()
+            self.assertTrue(title in saved_content)
+        finally:
+            if os.path.exists(temp_file):
+                os.remove(temp_file)