Create output directory for debugging images if it doesn't exist already.

alimuldal · alimuldal · commit 698d3627a80d · 2018-10-10T20:25:39.000+01:00
Also log the output paths in the error message for convenience.

PiperOrigin-RevId: 188318258
diff --git a/dm_control/mujoco/testing/image_utils.py b/dm_control/mujoco/testing/image_utils.py
@@ -22,6 +22,7 @@
 import collections
 import functools
 import os
+import sys
 
 # Internal dependencies.
 from dm_control import mujoco
@@ -194,17 +195,19 @@ def decorated_method(*args, **kwargs):
       try:
         test_method(*args, **kwargs)
       except ImagesNotClose as e:
+        _, _, tb = sys.exc_info()
+        if not os.path.exists(output_dir):
+          os.makedirs(output_dir)
         difference = e.actual.astype(np.double) - e.expected
         difference = (0.5 * (difference + 255)).astype(np.uint8)
-        _save_pixels(e.expected,
-                     os.path.join(output_dir,
-                                  '{}-expected.png'.format(method_name)))
-        _save_pixels(e.actual,
-                     os.path.join(output_dir,
-                                  '{}-actual.png'.format(method_name)))
-        _save_pixels(difference,
-                     os.path.join(output_dir,
-                                  '{}-difference.png'.format(method_name)))
-        raise  # Reraise the exception with the original traceback.
+        base_name = os.path.join(output_dir, method_name)
+        _save_pixels(e.expected, base_name + '-expected.png')
+        _save_pixels(e.actual, base_name + '-actual.png')
+        _save_pixels(difference, base_name + '-difference.png')
+        msg = ('{}. Debugging images saved to '
+               '{}-{{expected,actual,difference}}.png.'.format(e, base_name))
+        new_e = ImagesNotClose(msg, expected=e.expected, actual=e.actual)
+        # Reraise the exception with the original traceback.
+        six.reraise(ImagesNotClose, new_e, tb)
     return decorated_method
   return decorator
diff --git a/dm_control/mujoco/testing/image_utils_test.py b/dm_control/mujoco/testing/image_utils_test.py
@@ -61,7 +61,8 @@ def test_save_images_on_failure(self):
     def func():
       raise image_utils.ImagesNotClose(message, image1, image2)
 
-    with self.assertRaisesWithLiteralMatch(image_utils.ImagesNotClose, message):
+    with self.assertRaisesRegexp(image_utils.ImagesNotClose,
+                                 '{}.*'.format(message)):
       func()
 
     def validate_saved_file(name, expected_contents):
