diff --git a/docs/GenMultiDecay_Tutorial.ipynb b/docs/GenMultiDecay_Tutorial.ipynb
index 8bb26c39..8d88b6ed 100644
--- a/docs/GenMultiDecay_Tutorial.ipynb
+++ b/docs/GenMultiDecay_Tutorial.ipynb
@@ -13,8 +13,8 @@
     "# Tutorial for *GenMultiDecay* class\n",
     "This tutorial shows how ``phasespace.fromdecay.GenMultiDecay`` can be used.\n",
     "\n",
-    "In order to use this functionality, you need to install the extra `fromdecay`, for example through\n",
-    "``pip install phasespace[fromdecay]``.\n",
+    "In order to use this functionality, you need to install the extra dependencies, for example through\n",
+    "`pip install phasespace[fromdecay]`.\n",
     "\n",
     "This submodule makes it possible for `phasespace` and [`DecayLanguage`](https://github.com/scikit-hep/decaylanguage/) to work together.\n",
     "More generally, `GenMultiDecay` can also be used as a high-level interface for simulating particles that can decay in multiple different ways."
@@ -35,7 +35,7 @@
     "\n",
     "import zfit\n",
     "from particle import Particle\n",
-    "from decaylanguage import DecFileParser, DecayChainViewer\n",
+    "from decaylanguage import DecFileParser, DecayChainViewer, DecayChain, DecayMode\n",
     "import tensorflow as tf\n",
     "\n",
     "from phasespace.fromdecay import GenMultiDecay"
@@ -98,6 +98,25 @@
     "DecayChainViewer(pi0_chain)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can also create a decay using the `DecayChain` and `DecayMode` classes. However, a DecayChain can only contain one chain, i.e., a particle cannot decay in multiple ways."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dplus_decay = DecayMode(1, \"K- pi+ pi+ pi0\", model=\"PHSP\")\n",
+    "pi0_decay = DecayMode(1, \"gamma gamma\")\n",
+    "dplus_single = DecayChain(\"D+\", {\"D+\": dplus_decay, \"pi0\": pi0_decay})\n",
+    "DecayChainViewer(dplus_single.to_dict())"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -266,9 +285,45 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "If you want all $\\pi^0$ particles to decay with the same mass function, you do not need to specify the `zfit` parameter for each decay in the `dict`. Instead, one can pass the `particle_model_map` parameter to the constructor:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "GenMultiDecay.from_dict(dsplus_chain, particle_model_map={'pi0': 'gauss'})    # pi0 always decays with a gaussian mass distribution."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "When using `DecayChain`s, the syntax for specifying the mass function becomes cleaner:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dplus_decay = DecayMode(1, \"K- pi+ pi+ pi0\", model=\"PHSP\")  # The model parameter will be ignored by GenMultiDecay\n",
+    "pi0_decay = DecayMode(1, \"gamma gamma\", zfit=\"gauss\")   # Make pi0 have a gaussian mass distribution\n",
+    "dplus_single = DecayChain(\"D+\", {\"D+\": dplus_decay, \"pi0\": pi0_decay})\n",
+    "GenMultiDecay.from_dict(dplus_single.to_dict())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Custom mass functions\n",
     "The built-in supported mass function names are `gauss`, `bw`, and `relbw`, with `gauss` being the gaussian distribution, `bw` being the Breit-Wigner distribution, and `relbw` being the relativistic Breit-Wigner distribution.\n",
     "\n",
-    "If a non-supported value for the `zfit` parameter is used or if it is not specified, it will automatically use the relativistic Breit-Wigner distribution. This behavior can be changed by changing the value of `GenMultiDecay.DEFAULT_MASS_FUNC` to a different string, e.g., `\"gauss\"`.\n",
+    "If a non-supported value for the `zfit` parameter is not specified, it will automatically use the relativistic Breit-Wigner distribution. This behavior can be changed by changing the value of `GenMultiDecay.DEFAULT_MASS_FUNC` to a different string, e.g., `\"gauss\"`. If an invalid value for the `zfit` parameter is used, a `KeyError` is raised.\n",
     "\n",
     "It is also possible to add your own mass functions besides the built-in ones. You should then create a function that takes the mass and width of a particle and returns a mass function which with the [format](https://phasespace.readthedocs.io/en/stable/usage.html#resonances-with-variable-mass) that is used for all phasespace mass functions. Below is an example of a custom gaussian distribution (implemented in the same way as the built-in gaussian distribution), which uses `zfit` PDFs:"
    ]
@@ -362,4 +417,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 1
-}
\ No newline at end of file
+}
diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py
index 474dff91..081edbf5 100644
--- a/phasespace/fromdecay/genmultidecay.py
+++ b/phasespace/fromdecay/genmultidecay.py
@@ -2,7 +2,6 @@
 
 import itertools
 from collections.abc import Callable
-from typing import Union
 
 import tensorflow as tf
 import tensorflow.experimental.numpy as tnp
@@ -31,7 +30,8 @@ def from_dict(
         cls,
         dec_dict: dict,
         mass_converter: dict[str, Callable] = None,
-        tolerance: float = MASS_WIDTH_TOLERANCE,
+        tolerance: float = None,
+        particle_model_map: dict[str, str] = None,
     ):
         """Create a `GenMultiDecay` instance from a dict in the `DecayLanguage` package format.
 
@@ -45,9 +45,14 @@ def from_dict(
                 This dict will be combined with the predefined mass functions in this package.
                 See the Example below or the tutorial for how to use this parameter.
             tolerance: Minimum mass width of the particle to use a mass function instead of
-                assuming the mass to be constant. The default value is defined by the class variable
-                MASS_WIDTH_TOLERANCE and can be customized if desired.
-
+                assuming the mass to be constant. If None, the default value, defined by the class variable
+                MASS_WIDTH_TOLERANCE, will be used. This value can be customized if desired.
+            particle_model_map: A dict where the key is a particle name and the value is a mass function name.
+                If a particle is specified in the particle_model_map, then all appearances of this particle
+                in dec_dict will get the same mass function. This way, one does not have to manually add the
+                zfit parameter to every place where this particle appears in dec_dict. If the zfit parameter
+                is specified for a particle which is also included in particle_model_map, the zfit parameter
+                mass function name will be prioritized.
         Returns:
             The created GenMultiDecay object.
 
@@ -70,8 +75,13 @@ def from_dict(
                 {'bf': 0.016,
                  'fs': ['D+', 'gamma']}]}
 
-            If the D0 particle should have a mass distribution of a gaussian when it decays into K- and pi+
-            a `zfit` parameter can be added to its decay dict:
+            If the D0 particle should have a mass distribution of a gaussian when it decays, one can pass the
+            `particle_model_map` parameter to `from_dict`:
+            >>> dst_gen = GenMultiDecay.from_dict(dst_chain, particle_model_map={"D0": "gauss"})
+            This will then set the mass function of D0 to a gaussian for all its decays.
+
+            If more custom control is required, e.g., if D0 can decay in multiple ways and one of the decays
+            should have a specific mass function, a `zfit` parameter can be added to its decay dict:
             >>> dst_chain["D*+"][0]["fs"][0]["D0"][0]["zfit"] = "gauss"
             >>> dst_chain
             {'D*+': [{'bf': 0.984,
@@ -81,11 +91,13 @@ def from_dict(
                     'pi+']},
                 {'bf': 0.016,
                 'fs': ['D+', 'gamma']}]}
-
             This dict can then be passed to `GenMultiDecay.from_dict`:
             >>> dst_gen = GenMultiDecay.from_dict(dst_chain)
+            This will now convert make the D0 particle have a gaussian mass function, only when it decays into
+            K- and pi+. In this case, there are no other ways that D0 can decay, so using `particle_model_map`
+            is a cleaner and easier option.
 
-            If the decay of the D0 particle instead should be modelled by a mass distribution that does not
+            If the decay of the D0 particle should be modelled by a mass distribution that does not
             come with the package, a custom distribution can be created:
             >>> def custom_gauss(mass, width):
             >>>     particle_mass = tf.cast(mass, tf.float64)
@@ -113,6 +125,10 @@ def from_dict(
             For a more in-depth tutorial, see the tutorial on GenMultiDecay in the
             [documentation](https://phasespace.readthedocs.io/en/stable/GenMultiDecay_Tutorial.html).
         """
+        if tolerance is None:
+            tolerance = cls.MASS_WIDTH_TOLERANCE
+        if particle_model_map is None:
+            particle_model_map = {}
         if mass_converter is None:
             total_mass_converter = DEFAULT_CONVERTER
         else:
@@ -120,16 +136,18 @@ def from_dict(
             total_mass_converter = {**DEFAULT_CONVERTER, **mass_converter}
 
         gen_particles = _recursively_traverse(
-            dec_dict, total_mass_converter, tolerance=tolerance
+            dec_dict,
+            total_mass_converter,
+            particle_model_map=particle_model_map,
+            tolerance=tolerance,
         )
         return cls(gen_particles)
 
     def generate(
         self, n_events: int, normalize_weights: bool = True, **kwargs
-    ) -> (
-        tuple[list[tf.Tensor], list[tf.Tensor]]
-        | tuple[list[tf.Tensor], list[tf.Tensor], list[tf.Tensor]]
-    ):
+    ) -> tuple[list[tf.Tensor], list[tf.Tensor]] | tuple[
+        list[tf.Tensor], list[tf.Tensor], list[tf.Tensor]
+    ]:
         """Generate four-momentum vectors from the decay(s).
 
         Args:
@@ -197,13 +215,15 @@ def _get_particle_mass(
     name: str,
     mass_converter: dict[str, Callable],
     mass_func: str,
-    tolerance: float = GenMultiDecay.MASS_WIDTH_TOLERANCE,
+    tolerance: float,
 ) -> Callable | float:
     """Get mass or mass function of particle using the particle package.
 
     Args:
         name: Name of the particle. Name must be recognizable by the particle package.
-        tolerance : See _recursively_traverse
+        mass_converter: See _recursively_traverse
+        mass_func: See the name of the mass function, e.g., 'rel-bw'. Must be a valid key for mass_converter.
+        tolerance: See _recursively_traverse
 
     Returns:
         A function if the mass has a width smaller than tolerance. Otherwise, return a constant mass.
@@ -220,19 +240,25 @@ def _get_particle_mass(
 def _recursively_traverse(
     decaychain: dict,
     mass_converter: dict[str, Callable],
+    particle_model_map: dict[str, str],
+    tolerance: float,
     preexisting_particles: set[str] = None,
-    tolerance: float = GenMultiDecay.MASS_WIDTH_TOLERANCE,
 ) -> list[tuple[float, GenParticle]]:
     """Create all possible GenParticles by recursively traversing a dict from DecayLanguage, see Examples.
 
     Args:
         decaychain: Decay chain with the format from DecayLanguage
+        mass_converter: Maps from mass function names to the actual callable mass function.
+        particle_model_map: See GenMultiDecay.from_dict.
         preexisting_particles: Names of all particles that have already been created.
         tolerance: Minimum mass width for a particle to set a non-constant mass to a particle.
+        If None, set to default value, given by GenMultiDecay.MASS_WIDTH_TOLERANCE
 
     Returns:
         The generated GenParticle instances, one for each possible way of the decay.
     """
+    if tolerance is None:
+        tolerance = GenMultiDecay.MASS_WIDTH_TOLERANCE
     # Get the only key inside the decaychain dict
     (original_mother_name,) = decaychain.keys()
 
@@ -264,8 +290,9 @@ def _recursively_traverse(
                 daughter = _recursively_traverse(
                     daughter_name,
                     mass_converter,
+                    particle_model_map,
+                    tolerance,
                     preexisting_particles,
-                    tolerance=tolerance,
                 )
             else:
                 raise TypeError(
@@ -279,10 +306,17 @@ def _recursively_traverse(
             if is_top_particle:
                 mother_mass = Particle.from_evtgen_name(original_mother_name).mass
             else:
+                if "zfit" in dm:
+                    mass_func = dm["zfit"]
+                elif original_mother_name in particle_model_map:
+                    mass_func = particle_model_map[original_mother_name]
+                else:
+                    mass_func = GenMultiDecay.DEFAULT_MASS_FUNC
+
                 mother_mass = _get_particle_mass(
                     original_mother_name,
                     mass_converter=mass_converter,
-                    mass_func=dm.get("zfit", GenMultiDecay.DEFAULT_MASS_FUNC),
+                    mass_func=mass_func,
                     tolerance=tolerance,
                 )
 
diff --git a/phasespace/fromdecay/mass_functions.py b/phasespace/fromdecay/mass_functions.py
index f4ea260b..6767414b 100644
--- a/phasespace/fromdecay/mass_functions.py
+++ b/phasespace/fromdecay/mass_functions.py
@@ -2,14 +2,15 @@
 import zfit
 import zfit_physics as zphys
 
-
 # TODO refactor these mass functions using e.g. a decorator.
 #  Right now there is a lot of code repetition.
-def gauss(mass, width):
+
+
+def gauss_factory(mass, width):
     particle_mass = tf.cast(mass, tf.float64)
     particle_width = tf.cast(width, tf.float64)
 
-    def mass_func(min_mass, max_mass, n_events):
+    def gauss(min_mass, max_mass, n_events):
         min_mass = tf.cast(min_mass, tf.float64)
         max_mass = tf.cast(max_mass, tf.float64)
         pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs="")
@@ -18,14 +19,14 @@ def mass_func(min_mass, max_mass, n_events):
             lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator
         )
 
-    return mass_func
+    return gauss
 
 
-def breitwigner(mass, width):
+def breitwigner_factory(mass, width):
     particle_mass = tf.cast(mass, tf.float64)
     particle_width = tf.cast(width, tf.float64)
 
-    def mass_func(min_mass, max_mass, n_events):
+    def bw(min_mass, max_mass, n_events):
         min_mass = tf.cast(min_mass, tf.float64)
         max_mass = tf.cast(max_mass, tf.float64)
         pdf = zfit.pdf.Cauchy(m=particle_mass, gamma=particle_width, obs="")
@@ -34,14 +35,14 @@ def mass_func(min_mass, max_mass, n_events):
             lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator
         )
 
-    return mass_func
+    return bw
 
 
-def relativistic_breitwigner(mass, width):
+def relativistic_breitwigner_factory(mass, width):
     particle_mass = tf.cast(mass, tf.float64)
     particle_width = tf.cast(width, tf.float64)
 
-    def mass_func(min_mass, max_mass, n_events):
+    def relbw(min_mass, max_mass, n_events):
         min_mass = tf.cast(min_mass, tf.float64)
         max_mass = tf.cast(max_mass, tf.float64)
         pdf = zphys.pdf.RelativisticBreitWigner(
@@ -54,11 +55,11 @@ def mass_func(min_mass, max_mass, n_events):
             lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator
         )
 
-    return mass_func
+    return relbw
 
 
 DEFAULT_CONVERTER = {
-    "gauss": gauss,
-    "bw": breitwigner,
-    "relbw": relativistic_breitwigner,
+    "gauss": gauss_factory,
+    "bw": breitwigner_factory,
+    "relbw": relativistic_breitwigner_factory,
 }
diff --git a/phasespace/phasespace.py b/phasespace/phasespace.py
index e5ddd291..68cfa8dc 100644
--- a/phasespace/phasespace.py
+++ b/phasespace/phasespace.py
@@ -1,7 +1,3 @@
-from __future__ import annotations
-
-from collections.abc import Callable
-
 #!/usr/bin/env python3
 # =============================================================================
 # @file   phasespace.py
@@ -14,10 +10,12 @@
 
       F. James, Monte Carlo Phase Space, CERN 68-15 (1968)
 """
+from __future__ import annotations
+
 import inspect
 import warnings
+from collections.abc import Callable
 from math import pi
-from typing import Optional, Union
 
 import tensorflow as tf
 import tensorflow.experimental.numpy as tnp
diff --git a/tests/fromdecay/example_decay_chains.py b/tests/fromdecay/example_decay_chains.py
index c04d867d..934cfd65 100644
--- a/tests/fromdecay/example_decay_chains.py
+++ b/tests/fromdecay/example_decay_chains.py
@@ -9,7 +9,7 @@
 
 # D+ particle with only one way of decaying
 dplus_decay = DecayMode(1, "K- pi+ pi+ pi0", model="PHSP", zfit="relbw")
-pi0_decay = DecayMode(1, "gamma gamma")
+pi0_decay = DecayMode(1, "gamma gamma", zfit="relbw")
 dplus_single = DecayChain("D+", {"D+": dplus_decay, "pi0": pi0_decay}).to_dict()
 
 # pi0 particle that can decay in 4 possible ways
@@ -17,13 +17,6 @@
 
 # D+ particle that decays into 4 particles, out of which one particle in turn decays in 4 different ways.
 dplus_4grandbranches = dfp.build_decay_chains("D+")
-# Specify different mass functions for the different decays of pi0
-mass_functions = ["relbw", "bw", "gauss"]
-
-for mass_function, decay_mode in zip(
-    mass_functions, dplus_4grandbranches["D+"][0]["fs"][-1]["pi0"]
-):
-    decay_mode["zfit"] = mass_function
 
 # D*+ particle that has multiple children, grandchild particles, many of which can decay in multiple ways.
 dstarplus_big_decay = dfp.build_decay_chains("D*+")
diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py
index 61a2487d..63574308 100644
--- a/tests/fromdecay/test_fromdecay.py
+++ b/tests/fromdecay/test_fromdecay.py
@@ -1,5 +1,9 @@
 from __future__ import annotations
 
+from copy import deepcopy
+
+import pytest
+from decaylanguage import DecayChain, DecayMode
 from numpy.testing import assert_almost_equal
 
 from phasespace.fromdecay import GenMultiDecay
@@ -9,7 +13,7 @@
 
 
 def check_norm(full_decay: GenMultiDecay, **kwargs) -> list[tuple]:
-    """Helper function that checks whether the normalize_weights argument works for GenMultiDecay.generate.
+    """Helper function that tests whether the normalize_weights argument works for GenMultiDecay.generate.
 
     Args:
         full_decay: full_decay.generate will be called.
@@ -34,8 +38,31 @@ def check_norm(full_decay: GenMultiDecay, **kwargs) -> list[tuple]:
     return all_return_args
 
 
+def test_invalid_chain():
+    """Test that a ValueError is raised when a value in the fs key is not a str or dict."""
+    dm1 = DecayMode(1, "K- pi+ pi+ pi0", model="PHSP", zfit="rel-BW")
+    dm2 = DecayMode(1, "gamma gamma")
+    dc = DecayChain("D+", {"D+": dm1, "pi0": dm2}).to_dict()
+    dc["D+"][0]["fs"][0] = 1
+    with pytest.raises(TypeError):
+        GenMultiDecay.from_dict(dc)
+
+
+def test_invalid_mass_func_name():
+    """Test if an invalid mass function name as the zfit parameter raises a KeyError."""
+    dm1 = DecayMode(1, "K- pi+ pi+ pi0", model="PHSP")
+    dm2 = DecayMode(1, "gamma gamma", zfit="invalid name")
+    dc = DecayChain("D+", {"D+": dm1, "pi0": dm2}).to_dict()
+    with pytest.raises(KeyError):
+        GenMultiDecay.from_dict(dc, tolerance=1e-10)
+
+
 def test_single_chain():
-    """Test converting a DecayLanguage dict with only one possible decay."""
+    """Test converting a DecayLanguage dict with only one possible decay.
+
+    Since dplus_single is constructed using DecayChain.to_dict, this also tests that the code works dicts
+    created from DecayChains, not just .dec files.
+    """
     container = GenMultiDecay.from_dict(
         example_decay_chains.dplus_single, tolerance=1e-10
     )
@@ -48,6 +75,7 @@ def test_single_chain():
     for p in gen.children:
         if "pi0" == p.name[:3]:
             assert not p.has_fixed_mass
+            assert p._mass.__name__ == "relbw"
         else:
             assert p.has_fixed_mass
 
@@ -73,18 +101,55 @@ def test_branching_children():
 
 def test_branching_grandchilden():
     """Test converting a DecayLanguage dict where children to the mother particle can decay in many ways."""
-    container = GenMultiDecay.from_dict(example_decay_chains.dplus_4grandbranches)
+    # Specify different mass functions for the different decays of pi0
+    decay_dict = deepcopy(example_decay_chains.dplus_4grandbranches)
+
+    # Add different zfit parameters to all pi0 decays. The fourth decay has no zfit parameter
+    for mass_function, decay_mode in zip(
+        ("relbw", "bw", "gauss"), decay_dict["D+"][0]["fs"][-1]["pi0"]
+    ):
+        decay_mode["zfit"] = mass_function
+
+    container = GenMultiDecay.from_dict(decay_dict, tolerance=1e-10)
+
     output_decays = container.gen_particles
     assert len(output_decays) == 4
     assert_almost_equal(sum(d[0] for d in output_decays), 1)
+
+    for p, mass_func in zip(
+        output_decays, ("relbw", "bw", "gauss", GenMultiDecay.DEFAULT_MASS_FUNC)
+    ):
+        gen_particle = p[1]  # Ignore probability
+        assert gen_particle.children[-1].name == "pi0"
+        # Check that the zfit parameter assigns the correct mass function
+        assert gen_particle.children[-1]._mass.__name__ == mass_func
+
+    check_norm(container, n_events=1)
+    check_norm(container, n_events=100)
+
+
+def test_particle_model_map():
+    """Test that the particle_model_map parameter works as intended."""
+    container = GenMultiDecay.from_dict(
+        example_decay_chains.dplus_4grandbranches,
+        particle_model_map={"pi0": "bw"},
+        tolerance=1e-10,
+    )
+    output_decays = container.gen_particles
+    assert len(output_decays) == 4
+    assert_almost_equal(sum(d[0] for d in output_decays), 1)
+    for p in output_decays:
+        gen_particle = p[1]  # Ignore probability
+        assert gen_particle.children[-1].name[:3] == "pi0"
+        # Check that particle_model_map has assigned the bw mass function to all pi0 decays.
+        assert gen_particle.children[-1]._mass.__name__ == "bw"
     check_norm(container, n_events=1)
     check_norm(container, n_events=100)
-    # TODO add more asserts here
 
 
 def test_mass_converter():
     """Test that the mass_converter parameter works as intended."""
-    dplus_4grandbranches_massfunc = example_decay_chains.dplus_4grandbranches.copy()
+    dplus_4grandbranches_massfunc = deepcopy(example_decay_chains.dplus_4grandbranches)
     dplus_4grandbranches_massfunc["D+"][0]["fs"][-1]["pi0"][-1]["zfit"] = "rel-BW"
     container = GenMultiDecay.from_dict(
         dplus_4grandbranches_massfunc,
@@ -105,9 +170,41 @@ def test_mass_converter():
 
 
 def test_big_decay():
+    """Create a GenMultiDecay object from a large dict with many branches and subbranches."""
     container = GenMultiDecay.from_dict(example_decay_chains.dstarplus_big_decay)
     output_decays = container.gen_particles
     assert_almost_equal(sum(d[0] for d in output_decays), 1)
     check_norm(container, n_events=1)
     check_norm(container, n_events=100)
     # TODO add more asserts here
+
+
+def test_mass_width_tolerance():
+    """Test changing the MASS_WIDTH_TOLERANCE class variable."""
+    GenMultiDecay.MASS_WIDTH_TOLERANCE = 1e-10
+    output_decays = GenMultiDecay.from_dict(
+        example_decay_chains.dplus_4grandbranches
+    ).gen_particles
+    for p in output_decays:
+        gen_particle = p[1]  # Ignore probability
+        assert gen_particle.children[-1].name[:3] == "pi0"
+        # Check that particle_model_map has assigned the bw mass function to all pi0 decays.
+        assert not gen_particle.children[-1].has_fixed_mass
+    # Restore class variable to not affect other tests
+    GenMultiDecay.MASS_WIDTH_TOLERANCE = 1e-10
+
+
+def test_default_mass_func():
+    """Test changing the DEFAULT_MASS_FUNC class variable."""
+    GenMultiDecay.DEFAULT_MASS_FUNC = "bw"
+    output_decays = GenMultiDecay.from_dict(
+        example_decay_chains.dplus_4grandbranches, tolerance=1e-10
+    ).gen_particles
+    for p in output_decays:
+        gen_particle = p[1]  # Ignore probability
+        assert gen_particle.children[-1].name[:3] == "pi0"
+        # Check that particle_model_map has assigned the bw mass function to all pi0 decays.
+        assert gen_particle.children[-1]._mass.__name__ == "bw"
+
+    # Restore class variable to not affect other tests
+    GenMultiDecay.DEFAULT_MASS_FUNC = "bw"
diff --git a/tests/fromdecay/test_mass_functions.py b/tests/fromdecay/test_mass_functions.py
index 339cbe44..e93e4c2f 100644
--- a/tests/fromdecay/test_mass_functions.py
+++ b/tests/fromdecay/test_mass_functions.py
@@ -40,7 +40,8 @@ def ref_mass_func(min_mass, max_mass, n_events):
 
 
 @pytest.mark.parametrize(
-    "function", (mf.gauss, mf.breitwigner, mf.relativistic_breitwigner)
+    "function",
+    (mf.gauss_factory, mf.breitwigner_factory, mf.relativistic_breitwigner_factory),
 )
 @pytest.mark.parametrize("size", (1, 10))
 def test_shape(function: Callable, size: int, params: tuple = (1.0, 1.0)):