fix test

Author: drammock

mne/viz/tests/test_raw.py

@@ -989,15 +989,11 @@ def test_plot_raw_filtered(filtorder, raw, browser_backend):
 
 def _check_ylabel_psd(ylabel, amplitude, dB, unit):
     """Check that the ylabel is correct."""
-    if amplitude:
-        assert "sqrt" in ylabel
-    else:
-        assert "Hz" in ylabel and "sqrt" not in ylabel
-    if dB:
-        assert "dB" in ylabel
-    else:
-        assert unit in ylabel, ylabel
-        assert "dB" not in ylabel
+    numerator = "dB" if dB else unit
+    denominator = r"\sqrt{\mathrm{Hz}}" if amplitude else r"\mathrm{Hz}"
+    reference_val = rf"\ \mathrm{{re}}\ 1\ \mathrm{{{unit}}}" if dB else ""
+    pattern = rf"\mathrm{{{numerator}}}/{denominator}{reference_val}"
+    assert pattern in ylabel, ylabel
 
 
 def test_plot_raw_psd(raw, raw_orig):
@@ -1051,18 +1047,18 @@ def test_plot_raw_psd(raw, raw_orig):
     with pytest.warns(UserWarning, match="[Infinite|Zero]"):
         spectrum = raw.compute_psd()
     for dB, amplitude in itertools.product((True, False), (True, False)):
-        with pytest.warns(UserWarning, match="[Infinite|Zero]"):
+        with pytest.warns(UserWarning, match="(Infinite|Zero)"):
             fig = spectrum.plot(average=True, dB=dB, amplitude=amplitude)
         # check grad axes
         title = fig.axes[0].get_title()
         ylabel = fig.axes[0].get_ylabel()
-        unit = r"fT/cm/\sqrt{Hz}" if amplitude else "(fT/cm)²/Hz"
+        unit = "fT/cm" if amplitude else "(fT/cm)^2"
         assert title == "Gradiometers", title
         _check_ylabel_psd(ylabel, amplitude, dB, unit)
         # check mag axes
         title = fig.axes[1].get_title()
         ylabel = fig.axes[1].get_ylabel()
-        unit = r"fT/\sqrt{Hz}" if amplitude else "fT²/Hz"
+        unit = "fT" if amplitude else "fT^2"
         assert title == "Magnetometers", title
         _check_ylabel_psd(ylabel, amplitude, dB, unit)