Updating a unit test.

kaanaksit · Jul 10, 2024 · 7360a8e · 7360a8e
1 parent df3ba57
commit 7360a8e
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/test/test_ColorVideoVDP.py b/test/test_ColorVideoVDP.py
@@ -5,14 +5,18 @@
 def test():
     try:
         import pycvvdp
-        print('ColorVideoVDP is imported.')
-        main()
-    except:
+    except ImportError:
         print('ColorVideoVDP is missing, consider installing by visiting: https://github.com/gfxdisp/ColorVideoVDP')
         assert True == True
+    if 'pycvvdp' in sys.modules:
+        import pycvvdp
+        print('ColorVideoVDP is imported.')
+        main(pycvvdp)
+    assert True == True
 
 
 def main(
+         pycvvdp,
          output_directory = 'test_output', 
          device = torch.device('cpu'),
          pixels_per_degree = 60,
@@ -43,8 +47,7 @@ def main(
                              device = device
                             )
     loss = colorvdp.loss(image_noisy, image, dim_order = 'CHW')
-    loss_value, dictionary = loss
-    print(loss_value)
+    print(loss)
     assert True == True