update dispatch tests for more coverage

Author: LarsKue

tests/test_utils/test_dispatch.py

@@ -1,201 +1,170 @@
 import keras
 import pytest
 
-# Import the dispatch functions
-from bayesflow.utils import find_network, find_permutation, find_pooling, find_recurrent_net
-from tests.utils import assert_allclose
+from bayesflow.utils import find_inference_network, find_distribution, find_summary_network
 
-# --- Tests for find_network.py ---
 
+# --- Tests for find_inference_network.py ---
 
-class DummyMLP:
-    def __init__(self, *args, **kwargs):
-        self.args = args
-        self.kwargs = kwargs
 
+class DummyInferenceNetwork:
+    def __init__(self, *a, **kw):
+        self.args = a
+        self.kwargs = kw
 
-def test_find_network_with_string(monkeypatch):
-    # Monkeypatch the MLP entry in bayesflow.networks
-    monkeypatch.setattr("bayesflow.networks.MLP", DummyMLP)
-
-    net = find_network("mlp", 1, key="value")
-    assert isinstance(net, DummyMLP)
-    assert net.args == (1,)
-    assert net.kwargs == {"key": "value"}
 
+@pytest.mark.parametrize(
+    "name,expected_class_path",
+    [
+        ("coupling_flow", "bayesflow.networks.CouplingFlow"),
+        ("flow_matching", "bayesflow.networks.FlowMatching"),
+        ("consistency_model", "bayesflow.networks.ConsistencyModel"),
+    ],
+)
+def test_find_inference_network_by_name(monkeypatch, name, expected_class_path):
+    # patch the expected class in bayesflow.networks
+    components = expected_class_path.split(".")
+    module_path = ".".join(components[:-1])
+    class_name = components[-1]
 
-def test_find_network_with_type():
-    class CustomNet:
-        def __init__(self, x):
-            self.x = x
+    dummy_cls = DummyInferenceNetwork
+    monkeypatch.setattr(f"{module_path}.{class_name}", dummy_cls)
 
-    net = find_network(CustomNet, 42)
-    assert isinstance(net, CustomNet)
-    assert net.x == 42
+    net = find_inference_network(name, 1, key="val")
+    assert isinstance(net, DummyInferenceNetwork)
+    assert net.args == (1,)
+    assert net.kwargs == {"key": "val"}
 
 
-def test_find_network_with_keras_layer():
+def test_find_inference_network_by_keras_layer():
     layer = keras.layers.Dense(10)
-    returned = find_network(layer)
-    assert returned is layer
-
-
-def test_find_network_invalid_type():
-    with pytest.raises(TypeError):
-        find_network(123)
+    result = find_inference_network(layer)
+    assert result is layer
 
 
-# --- Tests for find_permutation.py ---
+def test_find_inference_network_by_keras_model():
+    model = keras.models.Sequential()
+    result = find_inference_network(model)
+    assert result is model
 
 
-class DummyRandomPermutation:
-    def __init__(self, *args, **kwargs):
-        self.args = args
-        self.kwargs = kwargs
+def test_find_inference_network_unknown_name():
+    with pytest.raises(ValueError):
+        find_inference_network("unknown_network_name")
 
 
-class DummySwap:
-    def __init__(self, *args, **kwargs):
-        self.args = args
-        self.kwargs = kwargs
+def test_find_inference_network_invalid_type():
+    with pytest.raises(TypeError):
+        find_inference_network(12345)
 
 
-class DummyOrthogonalPermutation:
-    def __init__(self, *args, **kwargs):
-        self.args = args
-        self.kwargs = kwargs
+# --- Tests for find_distribution.py ---
 
 
-def test_find_permutation_random(monkeypatch):
-    type("dummy_mod", (), {"RandomPermutation": DummyRandomPermutation})
-    monkeypatch.setattr("bayesflow.networks.coupling_flow.permutations.RandomPermutation", DummyRandomPermutation)
-    perm = find_permutation("random", 99, flag=True)
-    assert isinstance(perm, DummyRandomPermutation)
-    assert perm.args == (99,)
-    assert perm.kwargs == {"flag": True}
+class DummyDistribution:
+    def __init__(self, *a, **kw):
+        self.args = a
+        self.kwargs = kw
 
 
 @pytest.mark.parametrize(
-    "name,dummy_cls",
-    [("swap", DummySwap), ("learnable", DummyOrthogonalPermutation), ("orthogonal", DummyOrthogonalPermutation)],
+    "name, expected_class_path",
+    [
+        ("normal", "bayesflow.distributions.DiagonalNormal"),
+        ("student", "bayesflow.distributions.DiagonalStudentT"),
+        ("student-t", "bayesflow.distributions.DiagonalStudentT"),
+        ("student_t", "bayesflow.distributions.DiagonalStudentT"),
+    ],
 )
-def test_find_permutation_by_name(monkeypatch, name, dummy_cls):
-    # Inject dummy classes for each permutation type
-    if name == "swap":
-        monkeypatch.setattr("bayesflow.networks.coupling_flow.permutations.Swap", dummy_cls)
-    else:
-        monkeypatch.setattr("bayesflow.networks.coupling_flow.permutations.OrthogonalPermutation", dummy_cls)
-    perm = find_permutation(name, "a", b="c")
-    assert isinstance(perm, dummy_cls)
-    assert perm.args == ("a",)
-    assert perm.kwargs == {"b": "c"}
-
+def test_find_distribution_by_name(monkeypatch, name, expected_class_path):
+    components = expected_class_path.split(".")
+    module_path = ".".join(components[:-1])
+    class_name = components[-1]
 
-def test_find_permutation_with_keras_layer():
-    layer = keras.layers.Activation("relu")
-    perm = find_permutation(layer)
-    assert perm is layer
+    dummy_cls = DummyDistribution
+    monkeypatch.setattr(f"{module_path}.{class_name}", dummy_cls)
 
+    dist = find_distribution(name, 10, a=5)
+    assert isinstance(dist, DummyDistribution)
+    assert dist.args == (10,)
+    assert dist.kwargs == {"a": 5}
 
-def test_find_permutation_with_none():
-    res = find_permutation(None)
-    assert res is None
-
-
-def test_find_permutation_invalid_type():
-    with pytest.raises(TypeError):
-        find_permutation(3.14)
 
+def test_find_distribution_none_returns_none():
+    assert find_distribution(None) is None
 
-# --- Tests for find_pooling.py ---
 
+def test_find_distribution_with_keras_layer():
+    layer = keras.layers.Dense(3)
+    result = find_distribution(layer)
+    assert result is layer
 
-def dummy_pooling_constructor(*args, **kwargs):
-    return {"args": args, "kwargs": kwargs}
 
+def test_find_distribution_mixture_raises():
+    with pytest.raises(ValueError):
+        find_distribution("mixture")
 
-def test_find_pooling_mean():
-    pooling = find_pooling("mean")
-    # Check that a keras Lambda layer is returned
-    assert isinstance(pooling, keras.layers.Lambda)
-    # Test that the lambda function produces a mean when applied to a sample tensor.
-
-    sample = keras.ops.convert_to_tensor([[1, 2], [3, 4]])
-    # Keras Lambda layers expect tensors via call(), here we simply call the layer's function.
-    result = pooling.call(sample)
-    assert_allclose(result, keras.ops.mean(sample, axis=-2))
-
-
-@pytest.mark.parametrize("name,func", [("max", keras.ops.max), ("min", keras.ops.min)])
-def test_find_pooling_max_min(name, func):
-    pooling = find_pooling(name)
-    assert isinstance(pooling, keras.layers.Lambda)
-
-    sample = keras.ops.convert_to_tensor([[1, 2], [3, 4]])
-    result = pooling.call(sample)
-    assert_allclose(result, func(sample, axis=-2))
-
-
-def test_find_pooling_learnable(monkeypatch):
-    # Monkey patch the PoolingByMultiHeadAttention in its module
-    class DummyPoolingAttention:
-        def __init__(self, *args, **kwargs):
-            self.args = args
-            self.kwargs = kwargs
-
-    monkeypatch.setattr("bayesflow.networks.transformers.pma.PoolingByMultiHeadAttention", DummyPoolingAttention)
-    pooling = find_pooling("learnable", 7, option="test")
-    assert isinstance(pooling, DummyPoolingAttention)
-    assert pooling.args == (7,)
-    assert pooling.kwargs == {"option": "test"}
 
+def test_find_distribution_invalid_name():
+    with pytest.raises(ValueError):
+        find_distribution("invalid_name")
 
-def test_find_pooling_with_constructor():
-    # Passing a type should result in an instance.
-    class DummyPooling:
-        def __init__(self, data):
-            self.data = data
 
-    pooling = find_pooling(DummyPooling, "dummy")
-    assert isinstance(pooling, DummyPooling)
-    assert pooling.data == "dummy"
+def test_find_distribution_invalid_type():
+    with pytest.raises(TypeError):
+        find_distribution(3.14)
 
 
-def test_find_pooling_with_keras_layer():
-    layer = keras.layers.ReLU()
-    pooling = find_pooling(layer)
-    assert pooling is layer
+# --- Tests for find_summary_network.py ---
 
 
-def test_find_pooling_invalid_type():
-    with pytest.raises(TypeError):
-        find_pooling(123)
+class DummySummaryNetwork:
+    def __init__(self, *a, **kw):
+        self.args = a
+        self.kwargs = kw
 
 
-# --- Tests for find_recurrent_net.py ---
+@pytest.mark.parametrize(
+    "name,expected_class_path",
+    [
+        ("deep_set", "bayesflow.networks.DeepSet"),
+        ("set_transformer", "bayesflow.networks.SetTransformer"),
+        ("fusion_transformer", "bayesflow.networks.FusionTransformer"),
+        ("time_series_transformer", "bayesflow.networks.TimeSeriesTransformer"),
+        ("time_series_network", "bayesflow.networks.TimeSeriesNetwork"),
+    ],
+)
+def test_find_summary_network_by_name(monkeypatch, name, expected_class_path):
+    components = expected_class_path.split(".")
+    module_path = ".".join(components[:-1])
+    class_name = components[-1]
 
+    dummy_cls = DummySummaryNetwork
+    monkeypatch.setattr(f"{module_path}.{class_name}", dummy_cls)
 
-def test_find_recurrent_net_lstm():
-    constructor = find_recurrent_net("lstm")
-    assert constructor is keras.layers.LSTM
+    net = find_summary_network(name, 22, flag=True)
+    assert isinstance(net, DummySummaryNetwork)
+    assert net.args == (22,)
+    assert net.kwargs == {"flag": True}
 
 
-def test_find_recurrent_net_gru():
-    constructor = find_recurrent_net("gru")
-    assert constructor is keras.layers.GRU
+def test_find_summary_network_by_keras_layer():
+    layer = keras.layers.Dense(1)
+    out = find_summary_network(layer)
+    assert out is layer
 
 
-def test_find_recurrent_net_with_keras_layer():
-    layer = keras.layers.SimpleRNN(5)
-    net = find_recurrent_net(layer)
-    assert net is layer
+def test_find_summary_network_by_keras_model():
+    model = keras.models.Sequential()
+    out = find_summary_network(model)
+    assert out is model
 
 
-def test_find_recurrent_net_invalid_name():
+def test_find_summary_network_unknown_name():
     with pytest.raises(ValueError):
-        find_recurrent_net("invalid_net")
+        find_summary_network("unknown_summary_net")
 
 
-def test_find_recurrent_net_invalid_type():
+def test_find_summary_network_invalid_type():
     with pytest.raises(TypeError):
-        find_recurrent_net(3.1415)
+        find_summary_network(0.1234)