Simplify testing and dummy data generation (#147)

This uses `tfds.testing.mock_data` to simplify our training integration tests.

Author: lgeiger

larq_zoo/training/data.py

@@ -49,8 +49,13 @@ class Default(ImageClassification):
     input_shape = Field((IMAGE_SIZE, IMAGE_SIZE, 3))
 
     def input(self, data, training):
+        image = data["image"]
+        # tfds.testing.mock_data currently doesn't correctly handle custom decoders.
+        # See https://github.com/tensorflow/datasets/pull/1861
+        if image.dtype.is_integer:
+            image = tf.image.encode_jpeg(image)
         return preprocess_image_bytes(
-            data["image"], is_training=training, image_size=IMAGE_SIZE
+            image, is_training=training, image_size=IMAGE_SIZE
         )
 
 

tests/dummy_datasets.py

@@ -1,23 +1,29 @@
+from unittest import mock
+
 import pytest
+import tensorflow_datasets as tfds
 from click.testing import CliRunner
+from zookeeper.tf.dataset import TFDSDataset
 
 from larq_zoo.training import basic_experiments
-from tests import dummy_datasets  # noqa
 
 
 @pytest.mark.parametrize("command", list(basic_experiments.cli.commands.keys()))
-def test_cli(command):
+@tfds.testing.mock_data(num_examples=2, data_dir="gs://tfds-data/dataset_info")
+@mock.patch.object(TFDSDataset, "num_examples", return_value=2)
+def test_cli(_, command):
     result = CliRunner().invoke(
         basic_experiments.cli,
         [
             command,
-            "dataset=DummyOxfordFlowers",
+            "dataset=ImageNet",
             "epochs=1",
             "batch_size=2",
             "validation_frequency=5",
             "--no-use_tensorboard",
             "--no-use_model_checkpointing",
             "--no-save_weights",
         ],
+        catch_exceptions=False,
     )
     assert result.exit_code == 0