diff --git a/README.md b/README.md
index 8a1e07a..7ffc7ee 100644
--- a/README.md
+++ b/README.md
@@ -7,6 +7,15 @@ A Python package for plotting the locations, directions or energy distributions
 pip install openmc_source_plotter
 ```
 
+temporary fix
+For fixed source sources it is currently necessary to use openmc version 0.11
+and also to point the ```openmc_exec``` path to the openmc executable
+This can be installed with:
+```bash
+conda install -c conda-forge openmc=0.11
+```
+
+
 # Features
 
 The package provides functions to:
@@ -33,15 +42,12 @@ my_source = openmc.Source()
 # sets the energy distribution to a Muir distribution neutrons for DT fusion neutrons
 my_source.energy = openmc.stats.Muir(e0=14080000.0, m_rat=5.0, kt=20000.0)
 
-# makes an initial_source.h5 file with details of the particles
-osp.create_initial_particles(
-    source=my_source,
-    number_of_particles=10000,
-)
-
 # plots the particle energy distribution
-plot = osp.plot_energy_from_initial_source(
-    energy_bins=np.linspace(0, 20e6, 100)
+plot = osp.plot_source_energy(
+    source=my_source,
+    number_of_particles=2000,
+    energy_bins=np.linspace(0, 20e6, 100),
+    openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
 )
 
 plot.show()
@@ -59,15 +65,13 @@ my_source = openmc.Source()
 # sets the direction to isotropic
 my_source.angle = openmc.stats.Isotropic()
 
-# makes an initial_source.h5 file with details of the particles
-initial_source_filename = osp.create_initial_particles(
+# plots the particle energy distribution
+plot = osp.plot_source_direction(
     source=my_source,
     number_of_particles=100,
+    openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
 )
 
-# plots the particle energy distribution
-plot = osp.plot_direction_from_initial_source(input_filename=initial_source_filename)
-
 plot.show()
 ```
 ![openmc particle source direction plot](https://user-images.githubusercontent.com/8583900/143615706-3b3a8467-0233-42d6-a66c-d536c80a01d8.png)
diff --git a/examples/example_plot_direction_from_initial_source.py b/examples/example_plot_direction_from_initial_source.py
deleted file mode 100644
index c46214b..0000000
--- a/examples/example_plot_direction_from_initial_source.py
+++ /dev/null
@@ -1,21 +0,0 @@
-import openmc_source_plotter as osp
-import openmc
-import numpy as np
-
-# initializes a new source object
-my_source = openmc.Source()
-
-# sets the direction to isotropic
-my_source.angle = openmc.stats.Isotropic()
-
-# makes an initial_source.h5 file with details of the particles
-initial_source_filename = osp.create_initial_particles(
-    source=my_source,
-    number_of_particles=100,
-    openmc_exec="/home/jshimwell/miniconda3/envs/openmc_0_11_0/bin/openmc",
-)
-
-# plots the particle energy distribution
-plot = osp.plot_direction_from_initial_source(input_filename=initial_source_filename)
-
-plot.show()
diff --git a/examples/example_plot_position_from_initial_plasma_source.py b/examples/example_plot_plasma_source_position.py
similarity index 78%
rename from examples/example_plot_position_from_initial_plasma_source.py
rename to examples/example_plot_plasma_source_position.py
index da2f42f..2a118d8 100644
--- a/examples/example_plot_position_from_initial_plasma_source.py
+++ b/examples/example_plot_plasma_source_position.py
@@ -1,6 +1,5 @@
 import openmc_source_plotter as osp
 from openmc_plasma_source import TokamakSource
-import openmc
 
 my_source = TokamakSource(
     elongation=1.557,
@@ -21,14 +20,12 @@
     ion_temperature_beta=6,
 ).make_openmc_sources()
 
-# makes an initial_source.h5 file with details of the particles
-initial_source_filename = osp.create_initial_particles(
+
+# plots the particle energy distribution
+plot = osp.plot_source_position(
     source=my_source,
     number_of_particles=10,
     openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
 )
 
-# plots the particle energy distribution
-plot = osp.plot_position_from_initial_source(input_filename=initial_source_filename)
-
 plot.show()
diff --git a/examples/example_plot_source_direction.py b/examples/example_plot_source_direction.py
new file mode 100644
index 0000000..d568f6f
--- /dev/null
+++ b/examples/example_plot_source_direction.py
@@ -0,0 +1,17 @@
+import openmc_source_plotter as osp
+import openmc
+
+# initializes a new source object
+my_source = openmc.Source()
+
+# sets the direction to isotropic
+my_source.angle = openmc.stats.Isotropic()
+
+# plots the particle energy distribution
+plot = osp.plot_source_direction(
+    source=my_source,
+    number_of_particles=100,
+    openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
+)
+
+plot.show()
diff --git a/examples/example_plot_energy_from_initial_source.py b/examples/example_plot_source_energy.py
similarity index 56%
rename from examples/example_plot_energy_from_initial_source.py
rename to examples/example_plot_source_energy.py
index 5757dac..0c5b221 100644
--- a/examples/example_plot_energy_from_initial_source.py
+++ b/examples/example_plot_source_energy.py
@@ -5,22 +5,15 @@
 # initialises a new source object
 my_source = openmc.Source()
 
-# sets the location of the source to x=0 y=0 z=0
-my_source.space = openmc.stats.Point((0, 0, 0))
-
 # sets the energy distribution to a Muir distribution neutrons
 my_source.energy = openmc.stats.Muir(e0=14080000.0, m_rat=5.0, kt=20000.0)
 
-# makes an initial_source.h5 file with details of the particles
-initial_source_filename = osp.create_initial_particles(
+# plots the particle energy distribution
+plot = osp.plot_source_energy(
     source=my_source,
     number_of_particles=10000,
+    energy_bins=np.linspace(0, 20e6, 100),
     openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
 )
 
-# plots the particle energy distribution
-plot = osp.plot_energy_from_initial_source(
-    energy_bins=np.linspace(0, 20e6, 100), input_filename=initial_source_filename
-)
-
 plot.show()
diff --git a/examples/example_plot_position_from_initial_source.py b/examples/example_plot_source_position.py
similarity index 67%
rename from examples/example_plot_position_from_initial_source.py
rename to examples/example_plot_source_position.py
index 76ee303..381c763 100644
--- a/examples/example_plot_position_from_initial_source.py
+++ b/examples/example_plot_source_position.py
@@ -1,6 +1,5 @@
 import openmc_source_plotter as osp
 import openmc
-import numpy as np
 
 # initialises a new source object
 my_source = openmc.Source()
@@ -19,15 +18,9 @@
     r=radius, phi=angle, z=z_values, origin=(0.0, 0.0, 0.0)
 )
 
-
-# makes an initial_source.h5 file with details of the particles
-initial_source_filename = osp.create_initial_particles(
-    source=my_source,
-    number_of_particles=10,
-    openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
-)
-
 # plots the particle energy distribution
-plot = osp.plot_position_from_initial_source(input_filename=initial_source_filename)
+plot = osp.plot_source_position(
+    source=my_source, openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc"
+)
 
 plot.show()
diff --git a/examples/example_plot_two_source_energies.py b/examples/example_plot_two_source_energies.py
new file mode 100644
index 0000000..2bd0dfd
--- /dev/null
+++ b/examples/example_plot_two_source_energies.py
@@ -0,0 +1,24 @@
+import openmc_source_plotter as osp
+import openmc
+import numpy as np
+
+# initialises a new source object
+my_dt_source = openmc.Source()
+
+# sets the energy distribution to a Muir distribution DT neutrons
+my_dt_source.energy = openmc.stats.Muir(e0=14080000.0, m_rat=5.0, kt=20000.0)
+
+# initialises a new source object
+my_dd_source = openmc.Source()
+# sets the energy distribution to a Muir distribution DD neutrons
+my_dd_source.energy = openmc.stats.Muir(e0=2080000.0, m_rat=2.0, kt=20000.0)
+
+# plots the particle energy distribution
+plot = osp.plot_source_energy(
+    source=[my_dt_source, my_dd_source],
+    number_of_particles=10000,
+    energy_bins=np.linspace(0, 20e6, 100),
+    openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
+)
+
+plot.show()
diff --git a/openmc_source_plotter/__init__.py b/openmc_source_plotter/__init__.py
index 9676cfd..50c4954 100644
--- a/openmc_source_plotter/__init__.py
+++ b/openmc_source_plotter/__init__.py
@@ -1,5 +1,5 @@
-from .core import plot_direction_from_initial_source
-from .core import plot_position_from_initial_source
-from .core import plot_energy_from_initial_source
-from .core import get_particle_data
-from .core import create_initial_particles
+from .utils import get_particle_data
+from .utils import create_initial_particles
+from .core import plot_source_direction
+from .core import plot_source_energy
+from .core import plot_source_position
diff --git a/openmc_source_plotter/core.py b/openmc_source_plotter/core.py
index cd44613..3643f6f 100644
--- a/openmc_source_plotter/core.py
+++ b/openmc_source_plotter/core.py
@@ -1,189 +1,174 @@
 #!/usr/bin/env python3
 
-"""Provides utilities for creating h5 files containing initial source
-information and then plotting that information"""
+"""Provides functions for plotting source information"""
+
+import tempfile
+from typing import List, Union
 
-import h5py
 import numpy as np
 import openmc
 import plotly.graph_objects as go
 
+from .utils import create_initial_particles, get_particle_data
 
-def create_initial_particles(
-    source: openmc.source, number_of_particles: int = 2000, openmc_exec="openmc"
-) -> str:
-    """Accepts an openmc source and creates an initial_source.h5 file that can
-    be used to find propties of the source particles such as initial x,y,z
-    coordinates, direction and energy.
 
-    Note, I've found it easiest to install the latest version of openmc as this
-    has the write_initial_source for fixed source problems and then in a
-    new empty conda environment I have installed openmc version 0.11.0 with the
-    'conda install -c conda-forge openmc==0.11.0' command. I then go back to
-    the conda environment with the latest openmc version and run the python
-    script while setting the path to openmc_0.11.0 as the openmc_exec argument.
-    In my case it is '/home/jshimwell/miniconda3/envs/openmc_0_11_0/bin/openmc'
+def plot_source_energy(
+    source: Union[openmc.Source, List[openmc.Source]],
+    number_of_particles: int = 2000,
+    openmc_exec="openmc",
+    energy_bins: np.array = np.linspace(0, 20e6, 50),
+):
+    """makes a plot of the energy distribution OpenMC source(s)
 
     Args:
-        source: the openmc source to use
-        number_of_particles: the number of particles to sample
-        openmc_exec: the path of openmc executable or executable name if it
-            appears in your system $PATH. Defaults to 'openmc' which will use
-            the default openmc in your system $PATH environmental variable.
-    Returns:
-        The filename of the initial source file created (initial_source.h5)
+        source: The openmc.Source object or list of openmc.Source objects to plot.
+        number_of_particles: The number of source samples to obtain, more will
+            take longer but produce a smoother plot.
+        energy_bins: A numpy array of energy bins to use as energy bins.
+        openmc_exec: The path of the openmc executable to use
     """
 
-    # no real materials are needed for finding the source
-    mats = openmc.Materials([])
-
-    # just a minimal geometry
-    outer_surface = openmc.Sphere(r=100000, boundary_type="vacuum")
-    cell = openmc.Cell(region=-outer_surface)
-    universe = openmc.Universe(cells=[cell])
-    geom = openmc.Geometry(universe)
-
-    # Instantiate a Settings object
-    settings = openmc.Settings()
-    settings.run_mode = "fixed source"
-    settings.particles = number_of_particles
-    settings.batches = 1
-    settings.inactive = 0
-    settings.write_initial_source = True
-    settings.source = source
-
-    model = openmc.model.Model(geom, mats, settings)
-
-    model.export_to_xml()
-
-    openmc.run(openmc_exec=openmc_exec)
-
-    return "initial_source.h5"
-
-
-def get_particle_data(input_filename: str):
-
-    f = h5py.File(input_filename, "r")
-    dset = f["source_bank"]
-
-    x_values = []
-    y_values = []
-    z_values = []
-    x_dir = []
-    y_dir = []
-    z_dir = []
-    e_values = []
-
-    for particle in dset:
-        x_values.append(particle[0][0])
-        y_values.append(particle[0][1])
-        z_values.append(particle[0][2])
-        x_dir.append(particle[1][0])
-        y_dir.append(particle[1][1])
-        z_dir.append(particle[1][2])
-        e_values.append(particle[2])
+    fig = go.Figure()
 
-    return {
-        "x_values": x_values,
-        "y_values": y_values,
-        "z_values": z_values,
-        "x_dir": x_dir,
-        "y_dir": y_dir,
-        "z_dir": z_dir,
-        "e_values": e_values,
-    }
+    if isinstance(source, openmc.Source):
+        source = [source]
 
+    for single_source in source:
+        tmp_filename = tempfile.mkstemp(suffix=".h5", prefix=f"openmc_source_")[1]
+        create_initial_particles(
+            source=single_source,
+            number_of_particles=number_of_particles,
+            openmc_exec=openmc_exec,
+            output_source_filename=tmp_filename,
+        )
 
-def plot_energy_from_initial_source(
-    energy_bins: np.array = np.linspace(0, 20e6, 50),
-    input_filename: str = "initial_source.h5",
-):
-    """makes a plot of the energy distribution of the source"""
+        print("getting particle data", tmp_filename)
+        data = get_particle_data(tmp_filename)
 
-    data = get_particle_data(input_filename)
+        e_values = data["e_values"]
 
-    e_values = data["e_values"]
+        # Calculate pdf for source energies
+        probability, bin_edges = np.histogram(e_values, energy_bins, density=True)
 
-    # Calculate pdf for source energies
-    probability, bin_edges = np.histogram(e_values, energy_bins, density=True)
-    fig = go.Figure()
-
-    # Plot source energy histogram
-    fig.add_trace(
-        go.Scatter(
-            x=energy_bins[:-1],
-            y=probability * np.diff(energy_bins),
-            line={"shape": "hv"},
-            hoverinfo="text",
-            name="particle direction",
+        # Plot source energy histogram
+        fig.add_trace(
+            go.Scatter(
+                x=energy_bins[:-1],
+                y=probability * np.diff(energy_bins),
+                line={"shape": "hv"},
+                hoverinfo="text",
+                name=tmp_filename,
+            )
         )
-    )
 
     fig.update_layout(
         title="Particle energy",
         xaxis={"title": "Energy (eV)"},
         yaxis={"title": "Probability"},
+        showlegend=True,
     )
 
     return fig
 
 
-def plot_position_from_initial_source(input_filename="initial_source.h5"):
-    """makes a plot of the initial creation locations of the particle source"""
-
-    data = get_particle_data(input_filename)
+def plot_source_position(
+    source: Union[openmc.Source, List[openmc.Source]],
+    number_of_particles: int = 2000,
+    openmc_exec="openmc",
+):
+    """makes a plot of the initial creation postions of an OpenMC source(s)
 
-    text = ["Energy = " + str(i) + " eV" for i in data["e_values"]]
+    Args:
+        source: The openmc.Source object or list of openmc.Source objects to plot.
+        number_of_particles: The number of source samples to obtain.
+        openmc_exec: The path of the openmc executable to use
+    """
 
     fig = go.Figure()
 
-    fig.add_trace(
-        go.Scatter3d(
-            x=data["x_values"],
-            y=data["y_values"],
-            z=data["z_values"],
-            hovertext=text,
-            text=text,
-            mode="markers",
-            marker={
-                "size": 2,
-                "color": data["e_values"],
-            },
+    if isinstance(source, openmc.Source):
+        source = [source]
+
+    for single_source in source:
+        tmp_filename = tempfile.mkstemp(suffix=".h5", prefix=f"openmc_source_")[1]
+        create_initial_particles(
+            source=single_source,
+            number_of_particles=number_of_particles,
+            openmc_exec=openmc_exec,
+            output_source_filename=tmp_filename,
+        )
+
+        data = get_particle_data(tmp_filename)
+
+        text = ["Energy = " + str(i) + " eV" for i in data["e_values"]]
+
+        fig.add_trace(
+            go.Scatter3d(
+                x=data["x_values"],
+                y=data["y_values"],
+                z=data["z_values"],
+                hovertext=text,
+                text=text,
+                mode="markers",
+                marker={
+                    "size": 2,
+                    "color": data["e_values"],
+                },
+            )
         )
-    )
 
     fig.update_layout(title="Particle production coordinates - coloured by energy")
 
     return fig
 
 
-def plot_direction_from_initial_source(input_filename="initial_source.h5"):
-    """makes a plot of the initial creation directions of the particle source"""
-
-    data = get_particle_data(input_filename)
+def plot_source_direction(
+    source: Union[openmc.Source, List[openmc.Source]],
+    number_of_particles: int = 2000,
+    openmc_exec="openmc",
+):
+    """makes a plot of the initial creation directions of the particle source
 
+    Args:
+        source: The openmc.Source object or list of openmc.Source objects to plot.
+        number_of_particles: The number of source samples to obtain.
+        openmc_exec: The path of the openmc executable to use
+    """
     fig = go.Figure()
 
-    fig.add_trace(
-        {
-            "type": "cone",
-            "cauto": False,
-            "x": data["x_values"],
-            "y": data["y_values"],
-            "z": data["z_values"],
-            "u": data["x_dir"],
-            "v": data["y_dir"],
-            "w": data["z_dir"],
-            "cmin": 0,
-            "cmax": 1,
-            "anchor": "tail",
-            "colorscale": "Viridis",
-            "hoverinfo": "u+v+w+norm",
-            "sizemode": "absolute",
-            "sizeref": 30,
-            "showscale": False,
-        }
-    )
+    if isinstance(source, openmc.Source):
+        source = [source]
+
+    for single_source in source:
+        tmp_filename = tempfile.mkstemp(suffix=".h5", prefix=f"openmc_source_")[1]
+        create_initial_particles(
+            source=single_source,
+            number_of_particles=number_of_particles,
+            openmc_exec=openmc_exec,
+            output_source_filename=tmp_filename,
+        )
+        data = get_particle_data(tmp_filename)
+
+        fig.add_trace(
+            {
+                "type": "cone",
+                "cauto": False,
+                "x": data["x_values"],
+                "y": data["y_values"],
+                "z": data["z_values"],
+                "u": data["x_dir"],
+                "v": data["y_dir"],
+                "w": data["z_dir"],
+                "cmin": 0,
+                "cmax": 1,
+                "anchor": "tail",
+                "colorscale": "Viridis",
+                "hoverinfo": "u+v+w+norm",
+                "sizemode": "absolute",
+                "sizeref": 30,
+                "showscale": False,
+            }
+        )
 
     fig.update_layout(title="Particle initial directions")
 
diff --git a/openmc_source_plotter/utils.py b/openmc_source_plotter/utils.py
new file mode 100644
index 0000000..eb8b199
--- /dev/null
+++ b/openmc_source_plotter/utils.py
@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+
+"""Provides utilities for creating h5 files containing initial source
+information"""
+
+import h5py
+import openmc
+import shutil
+
+
+def get_particle_data(input_filename: str):
+
+    f = h5py.File(input_filename, "r")
+    dset = f["source_bank"]
+
+    x_values = []
+    y_values = []
+    z_values = []
+    x_dir = []
+    y_dir = []
+    z_dir = []
+    e_values = []
+
+    for particle in dset:
+        x_values.append(particle[0][0])
+        y_values.append(particle[0][1])
+        z_values.append(particle[0][2])
+        x_dir.append(particle[1][0])
+        y_dir.append(particle[1][1])
+        z_dir.append(particle[1][2])
+        e_values.append(particle[2])
+
+    return {
+        "x_values": x_values,
+        "y_values": y_values,
+        "z_values": z_values,
+        "x_dir": x_dir,
+        "y_dir": y_dir,
+        "z_dir": z_dir,
+        "e_values": e_values,
+    }
+
+
+def create_initial_particles(
+    source: openmc.source,
+    number_of_particles: int = 2000,
+    openmc_exec="openmc",
+    output_source_filename="initial_source.h5",
+) -> str:
+    """Accepts an openmc source and creates an initial_source.h5 file that can
+    be used to find propties of the source particles such as initial x,y,z
+    coordinates, direction and energy.
+
+    Note, I've found it easiest to install the latest version of openmc as this
+    has the write_initial_source for fixed source problems and then in a
+    new empty conda environment I have installed openmc version 0.11.0 with the
+    'conda install -c conda-forge openmc==0.11.0' command. I then go back to
+    the conda environment with the latest openmc version and run the python
+    script while setting the path to openmc_0.11.0 as the openmc_exec argument.
+    In my case it is '/home/jshimwell/miniconda3/envs/openmc_0_11_0/bin/openmc'
+
+    Args:
+        source: the openmc source to use
+        number_of_particles: the number of particles to sample
+        openmc_exec: the path of openmc executable or executable name if it
+            appears in your system $PATH. Defaults to 'openmc' which will use
+            the default openmc in your system $PATH environmental variable.
+        output_source_filename: the filename of the initial source h5 file
+            produced. Note openmc will write 'initial_source.h5' but this will
+            be moved to the output_source_filename specified
+    Returns:
+        The filename of the initial source file created (initial_source.h5)
+    """
+
+    # no real materials are needed for finding the source
+    mats = openmc.Materials([])
+
+    # just a minimal geometry
+    outer_surface = openmc.Sphere(r=100000, boundary_type="vacuum")
+    cell = openmc.Cell(region=-outer_surface)
+    universe = openmc.Universe(cells=[cell])
+    geom = openmc.Geometry(universe)
+
+    # Instantiate a Settings object
+    settings = openmc.Settings()
+    settings.run_mode = "fixed source"
+    settings.particles = number_of_particles
+    settings.batches = 1
+    settings.inactive = 0
+    settings.write_initial_source = True
+    settings.source = source
+
+    model = openmc.model.Model(geom, mats, settings)
+
+    model.export_to_xml()
+
+    openmc.run(openmc_exec=openmc_exec)
+
+    if output_source_filename != "initial_source.h5":
+        shutil.move("initial_source.h5", output_source_filename)
+
+    return output_source_filename
diff --git a/tests/test_core.py b/tests/test_core.py
new file mode 100644
index 0000000..b2b17bf
--- /dev/null
+++ b/tests/test_core.py
@@ -0,0 +1,65 @@
+from numpy import isin
+import openmc_source_plotter as osp
+import openmc
+import unittest
+from pathlib import Path
+import numpy as np
+import plotly.graph_objects as go
+
+
+class TestUtils(unittest.TestCase):
+    def setUp(self):
+
+        self.openmc_exec_dict = {
+            "ci": "/opt/openmc/build/bin/openmc",
+            "laptop": "/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
+            "desktop": "/home/jshimwell/miniconda3/envs/openmc_0_11_0/bin/openmc",
+        }
+        self.current_computer = "laptop"
+
+        # initialises a new source object
+        my_source = openmc.Source()
+
+        # sets the location of the source to x=0 y=0 z=0
+        my_source.space = openmc.stats.Point((0, 0, 0))
+
+        # sets the direction to isotropic
+        my_source.angle = openmc.stats.Isotropic()
+
+        # sets the energy distribution to 100% 14MeV neutrons
+        my_source.energy = openmc.stats.Discrete([14e6], [1])
+
+        self.my_source = my_source
+
+        # makes an initial_source.h5 file with details of the particles
+        self.initial_source_filename = osp.create_initial_particles(
+            source=my_source,
+            number_of_particles=10,
+            openmc_exec=self.openmc_exec_dict[self.current_computer],
+        )
+
+    def test_energy_plot(self):
+
+        plot = osp.plot_source_energy(
+            source=self.my_source,
+            number_of_particles=10000,
+            energy_bins=np.linspace(0, 20e6, 100),
+            openmc_exec=self.openmc_exec_dict[self.current_computer],
+        )
+        assert isinstance(plot, go.Figure)
+
+    def test_position_plot(self):
+
+        plot = osp.plot_source_position(
+            source=self.my_source,
+            openmc_exec=self.openmc_exec_dict[self.current_computer],
+        )
+        assert isinstance(plot, go.Figure)
+
+    def test_direction_plot(self):
+        plot = osp.plot_source_direction(
+            source=self.my_source,
+            number_of_particles=100,
+            openmc_exec=self.openmc_exec_dict["laptop"],
+        )
+        assert isinstance(plot, go.Figure)
diff --git a/tests/test_get_particle_data.py b/tests/test_utils.py
similarity index 52%
rename from tests/test_get_particle_data.py
rename to tests/test_utils.py
index 93d31c6..2c8648d 100644
--- a/tests/test_get_particle_data.py
+++ b/tests/test_utils.py
@@ -1,11 +1,19 @@
 import openmc_source_plotter as osp
 import openmc
 import unittest
+from pathlib import Path
 
 
-class TestReactor(unittest.TestCase):
+class TestUtils(unittest.TestCase):
     def setUp(self):
 
+        self.openmc_exec_dict = {
+            "ci": "/opt/openmc/build/bin/openmc",
+            "laptop": "/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
+            "desktop": "/home/jshimwell/miniconda3/envs/openmc_0_11_0/bin/openmc",
+        }
+        self.current_computer = "laptop"
+
         # initialises a new source object
         my_source = openmc.Source()
 
@@ -18,12 +26,13 @@ def setUp(self):
         # sets the energy distribution to 100% 14MeV neutrons
         my_source.energy = openmc.stats.Discrete([14e6], [1])
 
+        self.my_source = my_source
+
         # makes an initial_source.h5 file with details of the particles
         self.initial_source_filename = osp.create_initial_particles(
             source=my_source,
             number_of_particles=10,
-            openmc_exec="/opt/openmc/build/bin/openmc"
-            # openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
+            openmc_exec=self.openmc_exec_dict[self.current_computer],
         )
 
     def test_keys(self):
@@ -41,3 +50,15 @@ def test_keys(self):
         data = osp.get_particle_data(self.initial_source_filename)
         for key in key_values:
             assert key in data.keys()
+
+    def test_initial_source_output_file(self):
+        initial_source_filename = osp.create_initial_particles(
+            source=self.my_source,
+            number_of_particles=10,
+            output_source_filename="new_initial_source.h5",
+            openmc_exec=self.openmc_exec_dict[self.current_computer],
+        )
+
+        assert initial_source_filename == "new_initial_source.h5"
+        assert Path("new_initial_source.h5").exists()
+        assert Path("initial_source.h5").exists() is False