diff --git a/tests/test_colors.py b/tests/test_colors.py
index 52957fd..5765370 100755
--- a/tests/test_colors.py
+++ b/tests/test_colors.py
@@ -2,16 +2,20 @@
 import unittest
 import unittest.mock
 import io
+import pathlib
 
 from pywal import colors
 
 
+CACHE_DIR = pathlib.Path("/tmp/wal")
+
+
 class TestGenColors(unittest.TestCase):
     """Test the gen_colors functions."""
 
     def test_gen_colors(self):
         """> Generate a colorscheme."""
-        result = colors.get("tests/test_files/test.jpg")
+        result = colors.get("tests/test_files/test.jpg", CACHE_DIR)
         self.assertEqual(len(result["colors"]["color0"]), 7)
 
     def test_gen_colors_fail(self):
@@ -25,7 +29,7 @@ class TestGenColors(unittest.TestCase):
         # it's output so that we can read it.
         message = "colors: Found cached colorscheme."
         with unittest.mock.patch('sys.stdout', new=io.StringIO()) as fake_out:
-            colors.get("tests/test_files/test.jpg")
+            colors.get("tests/test_files/test.jpg", CACHE_DIR)
             self.assertEqual(fake_out.getvalue().strip(), message)
 
     def test_color_import(self):
diff --git a/tests/test_image.py b/tests/test_image.py
index 1cf7ec6..195951d 100644
--- a/tests/test_image.py
+++ b/tests/test_image.py
@@ -14,7 +14,7 @@ class TestImage(unittest.TestCase):
     def test_get_img_dir(self):
         """> Validate image directory."""
         result = image.get("tests/test_files")
-        self.assertEqual(result.endswith(".jpg"), True)
+        self.assertEqual(result.endswith((".jpg", ".png")), True)
 
     def test_get_img_fail(self):
         """> Validate image file. (fail)"""