diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 95e9f28fc8af6eb0579aeffa2171dc825f2b213d..5da3c2d78a9017e2f1862a6c8f6cce2630684e9c 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -7,7 +7,7 @@ before_script:
   - apt-get install --yes build-essential cmake libfmt-dev pybind11-dev graphviz python3-pyqt5 xvfb xdg-utils lcov
   - python -m pip install --upgrade pip
   - python --version
-  - pip install pytest pytest-cov sphinx furo numpydoc pytest-xvfb pytest-qt setuptools_scm sphinx-qt-documentation pytest-xdist
+  - pip install pytest pytest-cov sphinx furo numpydoc pytest-xvfb pytest-qt setuptools_scm sphinx-qt-documentation pytest-xdist pytest-mpl
   # - export CXXFLAGS='--coverage'
   - pip install -ve .
   # Move file, but should be handled by installation
@@ -17,7 +17,7 @@ before_script:
 .run-test:
   stage: test
   script:
-    - pytest --cov=b_asic --cov-report=xml:cov.xml --cov-report=term --color=yes test
+    - pytest --cov=b_asic --cov-report=xml:cov.xml --cov-report=term --color=yes test --mpl
     # - lcov --capture --directory . --output-file coverage.info
     # - lcov --output-file coverage.info --extract coverage.info $PWD/src/'*' $PWD/b_asic/'*'
     # - lcov --list coverage.info
diff --git a/b_asic/schedule.py b/b_asic/schedule.py
index 3bc6489bad921e38edf8523036fc79985fbf1988..c9312cc58bfdc06984d6554f4f5c8aed14e1b327 100644
--- a/b_asic/schedule.py
+++ b/b_asic/schedule.py
@@ -563,9 +563,12 @@ class Schedule:
         )
 
     def plot_schedule(self) -> None:
+        self._get_figure().show()
+
+    def _get_figure(self):
         fig, ax = plt.subplots()
         self._plot_schedule(ax)
-        fig.show()
+        return fig
 
     def _repr_svg_(self):
         fig, ax = plt.subplots()
diff --git a/test/baseline/test__get_figure_no_execution_times.png b/test/baseline/test__get_figure_no_execution_times.png
new file mode 100644
index 0000000000000000000000000000000000000000..10b2a22210ca0ceb3525263a9a8a6544c6b4a8ae
Binary files /dev/null and b/test/baseline/test__get_figure_no_execution_times.png differ
diff --git a/test/test_schedule.py b/test/test_schedule.py
index b4a1ffa31faf13dcbc3e08ac4c4b3d70dc811438..6678cfeb2f75562ccdaa98623068661cbce5649a 100644
--- a/test/test_schedule.py
+++ b/test/test_schedule.py
@@ -467,3 +467,9 @@ class TestProcesses:
     def test__get_memory_variables_list(self, secondorder_iir_schedule):
         mvl = secondorder_iir_schedule._get_memory_variables_list()
         assert len(mvl) == 12
+
+
+class TestFigureGeneration:
+    @pytest.mark.mpl_image_compare(remove_text=True, style='mpl20')
+    def test__get_figure_no_execution_times(self, secondorder_iir_schedule):
+        return secondorder_iir_schedule._get_figure()