Fix data cache test

larq · Aug 8, 2019 · f1ba210 · f1ba210
1 parent ca2b082
commit f1ba210
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/zookeeper/data_test.py b/zookeeper/data_test.py
@@ -1,11 +1,19 @@
 import pytest
 from unittest import mock
-from zookeeper import data
+from zookeeper import data, Preprocessing
+
+
+class MockPreprocessing(Preprocessing):
+    def inputs(self, data):
+        return data["image"]
+
+    def outputs(self, data):
+        return data["label"]
 
 
 @mock.patch("os.makedirs")
 def test_cache_dir(os_makedirs):
-    dataset = data.Dataset("mnist", lambda x: x)
+    dataset = data.Dataset("mnist", MockPreprocessing)
     assert dataset.get_cache_path("train") == None
 
     dataset.cache_dir = "memory"