actually skip mamba tests if environment is unsuitable

Author: LarsKue

tests/test_wrappers/test_mamba.py

@@ -3,25 +3,37 @@
 import bayesflow as bf
 
 
-@pytest.mark.torch
-def test_mamba_summary(random_time_series, mamba_summary_network):
+def should_skip():
+    import keras
+
+    if keras.backend.backend() != "torch":
+        return True, "Mamba tests can only be run on PyTorch."
+
     import torch
 
     if not torch.cuda.is_available():
-        pytest.skip("This test requires a GPU environment.")
+        return True, "Mamba tests can only be run on GPU."
+
+    try:
+        import mamba_ssm  # noqa: F401
+    except ImportError:
+        return True, "Could not import mamba."
+
+    return False, None
+
+
+skip, reason = should_skip()
 
+
+@pytest.mark.skipif(skip, reason=reason)
+def test_mamba_summary(random_time_series, mamba_summary_network):
     out = mamba_summary_network(random_time_series)
     # Batch size 2, summary dim 4
     assert out.shape == (2, 4)
 
 
-@pytest.mark.torch
+@pytest.mark.skipif(skip, reason=reason)
 def test_mamba_trains(random_time_series, inference_network, mamba_summary_network):
-    import torch
-
-    if not torch.cuda.is_available():
-        pytest.skip("This test requires a GPU environment.")
-
     workflow = bf.BasicWorkflow(
         inference_network=inference_network,
         summary_network=mamba_summary_network,