Updated tests with better fiure handling

entropicalabs · Feb 14, 2024 · c8d78d9 · c8d78d9
1 parent 14bb72b
commit c8d78d9
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/openqaoa-qiskit/tests/test_result_object_qiskit.py b/src/openqaoa-qiskit/tests/test_result_object_qiskit.py
@@ -1,5 +1,6 @@
 import unittest
 import networkx as nx
+import matplotlib.pyplot as plt
 
 from openqaoa import QAOA
 from openqaoa.problems import MinimumVertexCover
@@ -20,4 +21,7 @@ def test_plot_probabilities(self):
         q_shot.compile(vc)
         q_shot.optimize()
 
-        q_shot.result.plot_probabilities()
+        fig, ax = q_shot.result.plot_probabilities()
+        self.assertEqual(type(fig), plt.Figure)
+        self.assertEqual(type(ax), plt.Axes)
+        plt.close(fig)