diff --git a/src/ttmask/cone.py b/src/ttmask/cone.py
index 88a5574..d017f59 100644
--- a/src/ttmask/cone.py
+++ b/src/ttmask/cone.py
@@ -1,5 +1,4 @@
 from pathlib import Path
-
 import numpy as np
 import einops
 import typer
@@ -9,16 +8,13 @@
 from .soft_edge import add_soft_edge
 from .box_setup import box_setup
 
-
-@cli.command(name='cone')
 def cone(
-    sidelength: int = typer.Option(...),
-    cone_height: float = typer.Option(...),
-    cone_base_diameter: float = typer.Option(...),
-    soft_edge_width: int = typer.Option(0),
-    pixel_size: float = typer.Option(1),
-    output: Path = typer.Option(Path("cone.mrc"))
-):
+    sidelength: int, 
+    cone_height: float, 
+    cone_base_diameter: float, 
+    soft_edge_width: int, 
+    pixel_size: float
+) -> np.ndarray:
     # establish our coordinate system and empty mask
     coordinates_centered, mask = box_setup(sidelength)
     # distances between each pixel and center :
@@ -55,4 +51,19 @@ def cone(
     mask = np.roll(mask, z_shift, axis=0)
     mask = add_soft_edge(mask, soft_edge_width)
 
-    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
+    return mask
+
+@cli.command(name='cone')
+def cone_cli(
+    sidelength: int = typer.Option(...),
+    cone_height: float = typer.Option(...),
+    cone_base_diameter: float = typer.Option(...),
+    soft_edge_width: int = typer.Option(0),
+    pixel_size: float = typer.Option(1),
+    output: Path = typer.Option(Path("cone.mrc")),
+):
+    mask = cone(sidelength, cone_height, cone_base_diameter, soft_edge_width, pixel_size)
+
+    # Save the mask to an MRC file
+    with mrcfile.new(output, overwrite=True) as mrc:
+        mrc.set_data(mask.astype(np.float32))
diff --git a/src/ttmask/cube.py b/src/ttmask/cube.py
index ced32b0..8bdb4c0 100644
--- a/src/ttmask/cube.py
+++ b/src/ttmask/cube.py
@@ -1,5 +1,4 @@
 from pathlib import Path
-
 import numpy as np
 import typer
 import mrcfile
@@ -8,17 +7,14 @@
 from ._cli import cli
 from .box_setup import box_setup
 
-
-@cli.command(name='cube')
 def cube(
-    sidelength: int = typer.Option(...),
-    cube_sidelength: float = typer.Option(...),
-    soft_edge_width: float = typer.Option(0),
-    pixel_size: float = typer.Option(1),
-    output: Path = typer.Option(Path("cube.mrc")),
-    wall_thickness: float = typer.Option(0),
-):
-    # establish our coordinate system and empty mask
+    sidelength: int, 
+    cube_sidelength: float, 
+    soft_edge_width: float, 
+    pixel_size: float,
+    wall_thickness: float
+) -> np.ndarray:
+     # establish our coordinate system and empty mask
     coordinates_centered, mask = box_setup(sidelength)
     #converting relative coordinates to xyz distances (i.e. not a negative number) :
     xyz_distances = np.abs(coordinates_centered)
@@ -36,5 +32,19 @@ def cube(
     #if requested, a soft edge is added to the mask
     mask = add_soft_edge(mask, soft_edge_width)
 
-    #output created with desired pixel size.
-    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
+    return mask
+
+@cli.command(name='cube')
+def cube_cli(
+    sidelength: int = typer.Option(...),
+    cube_sidelength: float = typer.Option(...),
+    soft_edge_width: float = typer.Option(0),
+    pixel_size: float = typer.Option(1),
+    output: Path = typer.Option(Path("cube.mrc")),
+    wall_thickness: float = typer.Option(0),
+):
+    mask = cube(sidelength, cube_sidelength, soft_edge_width, pixel_size)
+
+    # Save the mask to an MRC file
+    with mrcfile.new(output, overwrite=True) as mrc:
+        mrc.set_data(mask.astype(np.float32))
diff --git a/src/ttmask/cuboid.py b/src/ttmask/cuboid.py
index e088e3d..c8539f0 100644
--- a/src/ttmask/cuboid.py
+++ b/src/ttmask/cuboid.py
@@ -1,5 +1,4 @@
 from pathlib import Path
-
 import numpy as np
 import typer
 from typing import Tuple
@@ -10,16 +9,13 @@
 from .soft_edge import add_soft_edge
 from .box_setup import box_setup
 
-
-@cli.command(name='cuboid')
 def cuboid(
-    sidelength: int = typer.Option(...),
-    cuboid_sidelengths: Annotated[Tuple[float, float, float], typer.Option()] = (None, None, None),
-    soft_edge_width: float = typer.Option(0),
-    pixel_size: float = typer.Option(1),
-    output: str = typer.Option(Path("cuboid.mrc")),
-    wall_thickness: float = typer.Option(0),
-):
+    sidelength: int, 
+    cuboid_sidelengths: Tuple[float, float, float], 
+    soft_edge_width: float, 
+    pixel_size: float,
+    wall_thickness: float
+) -> np.ndarray:
     # establish our coordinate system and empty mask
     coordinates_centered, mask = box_setup(sidelength)
     #converting relative coordinates to xyz distances (i.e. not a negative number) :
@@ -37,6 +33,20 @@ def cuboid(
 
     # if requested, a soft edge is added to the mask
     mask = add_soft_edge(mask, soft_edge_width)
+    
+    return mask
+
+@cli.command(name='cuboid')
+def cuboid_cli(
+    sidelength: int = typer.Option(...),
+    cuboid_sidelengths: Annotated[Tuple[float, float, float], typer.Option()] = (None, None, None),
+    soft_edge_width: float = typer.Option(0),
+    pixel_size: float = typer.Option(1),
+    wall_thickness: float = typer.Option(0),
+    output: Path = typer.Option(Path("cuboid.mrc")),
+):
+    mask = cuboid(sidelength, cuboid_sidelengths, soft_edge_width, pixel_size,wall_thickness)
 
-    # output created with desired pixel size.
-    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
+    # Save the mask to an MRC file
+    with mrcfile.new(output, overwrite=True) as mrc:
+        mrc.set_data(mask.astype(np.float32))
diff --git a/src/ttmask/curved_surface.py b/src/ttmask/curved_surface.py
index b19edb4..f71939f 100644
--- a/src/ttmask/curved_surface.py
+++ b/src/ttmask/curved_surface.py
@@ -1,6 +1,4 @@
 from pathlib import Path
-
-
 import numpy as np
 import typer
 import mrcfile
@@ -9,15 +7,13 @@
 from .soft_edge import add_soft_edge
 from .box_setup import box_setup
 
-@cli.command(name='curved_surface')
 def curved_surface(
-    sidelength: int = typer.Option(...),
-    fit_sphere_diameter: float = typer.Option(...),
-    soft_edge_width: int = typer.Option(0),
-    pixel_size: float = typer.Option(1),
-    output: Path = typer.Option(Path("curved_surface.mrc")),
-    surface_thickness: float = typer.Option(...),
-):
+    sidelength: int, 
+    fit_sphere_diameter: float, 
+    soft_edge_width: int, 
+    pixel_size: float,
+    surface_thickness: float 
+) -> np.ndarray:
     sphere_radius = fit_sphere_diameter / 2
 
     # establish our coordinate system and empty mask
@@ -41,5 +37,19 @@ def curved_surface(
     # if requested, a soft edge is added to the mask
     mask = add_soft_edge(mask, soft_edge_width)
 
-    # output created with desired pixel size.
-    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
+    return mask
+
+@cli.command(name='curved_surface')
+def curved_surface_cli(
+    sidelength: int = typer.Option(...),
+    fit_sphere_diameter: float = typer.Option(...),
+    soft_edge_width: int = typer.Option(0),
+    pixel_size: float = typer.Option(1),
+    output: Path = typer.Option(Path("curved_surface.mrc")),
+    surface_thickness: float = typer.Option(...),
+):
+    mask = curved_surface(sidelength, fit_sphere_diameter, soft_edge_width, pixel_size, surface_thickness)
+
+    # Save the mask to an MRC file
+    with mrcfile.new(output, overwrite=True) as mrc:
+        mrc.set_data(mask.astype(np.float32))
diff --git a/src/ttmask/cylinder.py b/src/ttmask/cylinder.py
index 923c4d8..ea412fc 100644
--- a/src/ttmask/cylinder.py
+++ b/src/ttmask/cylinder.py
@@ -5,20 +5,18 @@
 import mrcfile
 
 from ._cli import cli
+
 from .soft_edge import add_soft_edge
 from .box_setup import box_setup
 
-
-@cli.command(name='cylinder')
 def cylinder(
-    sidelength: int = typer.Option(...),
-    cylinder_height: float = typer.Option(...),
-    cylinder_diameter: float = typer.Option(...),
-    wall_thickness: float = typer.Option(0),
-    soft_edge_width: int = typer.Option(0),
-    pixel_size: float = typer.Option(1),
-    output: Path = typer.Option(Path("cylinder.mrc"))
-):
+    sidelength: int, 
+    cylinder_height: float, 
+    cylinder_diameter: float, 
+    wall_thickness: float, 
+    soft_edge_width: int,
+    pixel_size: float
+) -> np.ndarray:
     cylinder_radius = cylinder_diameter / 2
 
     # establish our coordinate system and empty mask
@@ -43,5 +41,21 @@ def cylinder(
     # if requested, a soft edge is added to the mask
     mask = add_soft_edge(mask, soft_edge_width)
 
-    # output created with desired pixel size.
-    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
+    return mask
+
+
+@cli.command(name='cylinder')
+def cylinder_cli(
+    sidelength: int = typer.Option(...),
+    cylinder_height: float = typer.Option(...),
+    cylinder_diameter: float = typer.Option(...),
+    wall_thickness: float = typer.Option(0),
+    soft_edge_width: int = typer.Option(0),
+    pixel_size: float = typer.Option(1),
+    output: Path = typer.Option(Path("cylinder.mrc")),
+):
+    mask = cylinder(sidelength, cylinder_height, cylinder_diameter, wall_thickness, soft_edge_width, pixel_size)
+
+    # Save the mask to an MRC file
+    with mrcfile.new(output, overwrite=True) as mrc:
+        mrc.set_data(mask.astype(np.float32))
diff --git a/src/ttmask/ellipsoid.py b/src/ttmask/ellipsoid.py
index 0f65a76..0395d6c 100644
--- a/src/ttmask/ellipsoid.py
+++ b/src/ttmask/ellipsoid.py
@@ -10,17 +10,13 @@
 from ._cli import cli
 from .box_setup import box_setup
 
-
-@cli.command(name='ellipsoid')
 def ellipsoid(
-
-    sidelength: int = typer.Option(...),
-    ellipsoid_dimensions: Annotated[Tuple[float, float, float], typer.Option()] = (None, None, None),
-    soft_edge_width: int = typer.Option(0),
-    pixel_size: float = typer.Option(1),
-    output: Path = typer.Option(Path("ellipsoid.mrc")),
-    wall_thickness: float = typer.Option(0),
-):
+    sidelength: int, 
+    ellipsoid_dimensions: Tuple[float, float, float], 
+    soft_edge_width: float,
+    pixel_size: float,
+    wall_thickness: float
+) -> np.ndarray:
     # establish our coordinate system and empty mask
     coordinates_centered, mask = box_setup(sidelength)
     #converting relative coordinates to xyz distances (i.e. not a negative number) :
@@ -51,5 +47,18 @@ def ellipsoid(
     # if requested, a soft edge is added to the mask
     mask = add_soft_edge(mask, soft_edge_width)
 
-    # output created with desired pixel size.
-    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
+    return mask
+
+@cli.command(name='ellipsoid')
+def ellipsoid_cli(
+    sidelength: int = typer.Option(...),
+    ellipsoid_dimensions: Annotated[Tuple[float, float, float], typer.Option()] = (None, None, None),
+    soft_edge_width: float = typer.Option(0),
+    pixel_size: float = typer.Option(1),
+    output: Path = typer.Option(Path("ellipsoid.mrc")),
+):
+    mask = ellipsoid(sidelength, ellipsoid_dimensions, soft_edge_width, pixel_size)
+
+    # Save the mask to an MRC file
+    with mrcfile.new(output, overwrite=True) as mrc:
+        mrc.set_data(mask.astype(np.float32))
diff --git a/src/ttmask/map2mask.py b/src/ttmask/map2mask.py
index 90bec28..38f37fa 100644
--- a/src/ttmask/map2mask.py
+++ b/src/ttmask/map2mask.py
@@ -7,18 +7,12 @@
 from .soft_edge import add_soft_edge
 from .add_padding import add_padding
 
-@cli.command(name='map2mask')
-def map2mask(
-
-    input_map: Path = typer.Option(Path("map.mrc")),
-    binarization_threshold: float = typer.Option(...),
-    output_mask: Path = typer.Option(Path("mask.mrc")),
-    pixel_size: float = typer.Option(...),
-    soft_edge_width: int = typer.Option(0),
-    padding_width: int = typer.Option(0),
-):
-    with mrcfile.open(input_map) as mrc:
-        map_data = np.array(mrc.data)
+def mask_from_map(
+    map_data: np.ndarray, 
+    binarization_threshold: float, 
+    soft_edge_width: int, 
+    padding_width: int
+) -> np.ndarray:
 
     above_threshold = map_data >= binarization_threshold
     below_threshold = map_data < binarization_threshold
@@ -29,4 +23,21 @@ def map2mask(
     padded_mask = add_padding(map_data, padding_width)
     mask = add_soft_edge(padded_mask, soft_edge_width)
 
-    mrcfile.write(output_mask, mask, voxel_size=pixel_size, overwrite=True)
\ No newline at end of file
+    return mask
+
+@cli.command(name='map2mask')
+def map2mask(
+    input_map: Path = typer.Option(Path("map.mrc")),
+    binarization_threshold: float = typer.Option(...),
+    output_mask: Path = typer.Option(Path("mask.mrc")),
+    pixel_size: float = typer.Option(...),
+    soft_edge_width: int = typer.Option(0),
+    padding_width: int = typer.Option(0),
+):
+    with mrcfile.open(input_map, permissive=True) as mrc:
+        data = mrc.data
+    mask = mask_from_map(data, binarization_threshold, soft_edge_width, padding_width)
+
+    # Save the mask to an MRC file
+    with mrcfile.new(output_mask, overwrite=True) as mrc:
+        mrc.set_data(mask.astype(np.float32))
\ No newline at end of file
diff --git a/src/ttmask/sphere.py b/src/ttmask/sphere.py
index f7217e6..e59e577 100644
--- a/src/ttmask/sphere.py
+++ b/src/ttmask/sphere.py
@@ -1,6 +1,4 @@
 from pathlib import Path
-
-
 import numpy as np
 import typer
 import mrcfile
@@ -9,21 +7,19 @@
 from .soft_edge import add_soft_edge
 from .box_setup import box_setup
 
-@cli.command(name='sphere')
 def sphere(
-    sidelength: int = typer.Option(...),
-    sphere_diameter: float = typer.Option(...),
-    soft_edge_width: int = typer.Option(0),
-    pixel_size: float = typer.Option(1),
-    output: Path = typer.Option(Path("sphere.mrc")),
-    wall_thickness: float = typer.Option(0),
-):
+    sidelength: int, 
+    sphere_diameter: float, 
+    soft_edge_width: int, 
+    pixel_size: float, 
+    wall_thickness: float
+) -> np.ndarray:
     sphere_radius = sphere_diameter / 2
 
     # establish our coordinate system and empty mask
     coordinates_centered, mask = box_setup(sidelength)
 
-    #determine distances of each pixel to the center
+    # determine distances of each pixel to the center
     distance_to_center = np.linalg.norm(coordinates_centered, axis=-1)
 
     # set up criteria for which pixels are inside the sphere and modify values to 1.
@@ -38,5 +34,19 @@ def sphere(
     # if requested, a soft edge is added to the mask
     mask = add_soft_edge(mask, soft_edge_width)
 
-    # output created with desired pixel size.
-    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
+    return mask
+
+@cli.command(name='sphere')
+def sphere_cli(
+    sidelength: int = typer.Option(...),
+    sphere_diameter: float = typer.Option(...),
+    soft_edge_width: int = typer.Option(0),
+    pixel_size: float = typer.Option(1),
+    output: Path = typer.Option(Path("sphere.mrc")),
+    wall_thickness: float = typer.Option(0),
+):
+    mask = sphere(sidelength, sphere_diameter, soft_edge_width, pixel_size, wall_thickness)
+
+    # Save the mask to an MRC file
+    with mrcfile.new(output, overwrite=True) as mrc:
+        mrc.set_data(mask.astype(np.float32))
diff --git a/src/ttmask/tube.py b/src/ttmask/tube.py
index fbade89..db680d7 100644
--- a/src/ttmask/tube.py
+++ b/src/ttmask/tube.py
@@ -1,7 +1,6 @@
 from pathlib import Path
 
 import numpy as np
-
 import typer
 import mrcfile
 
@@ -9,17 +8,14 @@
 from .soft_edge import add_soft_edge
 from .box_setup import box_setup
 
-
-@cli.command(name='tube')
 def tube(
-    sidelength: int = typer.Option(...),
-    tube_height: float = typer.Option(...),
-    tube_diameter: float = typer.Option(...),
-    wall_thickness: float = typer.Option(0),
-    soft_edge_width: int = typer.Option(0),
-    pixel_size: float = typer.Option(1),
-    output: Path = typer.Option(Path("tube.mrc")),
-):
+    sidelength: int, 
+    tube_height: float, 
+    tube_diameter: float, 
+    wall_thickness: float,
+    soft_edge_width: int,
+    pixel_size: float,
+) -> np.ndarray:
     tube_radius = tube_diameter / 2
 
     # establish our coordinate system and empty mask
@@ -41,5 +37,18 @@ def tube(
     # if requested, a soft edge is added to the mask
     mask = add_soft_edge(mask, soft_edge_width)
 
-    # output created with desired pixel size.
-    mrcfile.write(output, mask, voxel_size=pixel_size, overwrite=True)
+    return mask
+
+@cli.command(name='tube')
+def tube_cli(
+    sidelength: int = typer.Option(...),
+    tube_height: float = typer.Option(...),
+    tube_diameter: float = typer.Option(...),
+    wall_thickness: float = typer.Option(0),
+    output: Path = typer.Option(Path("tube.mrc")),
+):
+    mask = tube(sidelength, tube_height, tube_diameter, wall_thickness)
+
+    # Save the mask to an MRC file
+    with mrcfile.new(output, overwrite=True) as mrc:
+        mrc.set_data(mask.astype(np.float32))
diff --git a/tests/test_functions.py b/tests/test_functions.py
new file mode 100644
index 0000000..d96c2d9
--- /dev/null
+++ b/tests/test_functions.py
@@ -0,0 +1,53 @@
+from ttmask import cone, cube, cuboid, cylinder, ellipsoid, map2mask, sphere, tube
+import numpy as np
+
+def test_cone():
+    mask = cone(100, 50, 50, 0, 1)
+    assert mask.shape == (100, 100, 100)
+    assert mask.sum() > np.pi * 24**2 * (50 / 3) # Volume of cone
+    assert mask.sum() < np.pi * 25**2 * 50 # Volume of cylinder
+
+def test_cube():
+    mask = cube(100, 50, 0, 1, 0)
+    assert mask.shape == (100, 100, 100)
+    # Test against volume of cube +- center and subpixel issues
+    assert mask.sum() > 48**3
+    assert mask.sum() < 52**3
+
+def test_cuboid():
+    mask = cuboid(100, (50,40,30), 0, 1, 0)
+    assert mask.shape == (100, 100, 100)
+    # Test against volume of cuboid +- center and subpixel issues
+    assert mask.sum() > 48 * 38 * 28
+    assert mask.sum() < 52 * 42 * 32
+
+#def test_curved_surface():
+#    mask = curved_surface(100, 50, 50, 0, 1)
+#    assert mask.shape == (100, 100, 100)
+#    assert mask.sum() > 2 * np.pi * 25**2 # Area of cylinder
+#    assert mask.sum() < 2 * np.pi * 25 * 50 # Area of cylinder
+
+def test_cylinder():
+    mask = cylinder(100, 50, 50, 0, 0, 1)
+    assert mask.shape == (100, 100, 100)
+    assert mask.sum() > np.pi * 25**2 * 48 # Volume of cylinder
+    assert mask.sum() < np.pi * 25**2 * 51 # Volume of cylinder
+
+def test_ellipsoid():
+    mask = ellipsoid(100, (50,40,30), 0, 1,0)
+    assert mask.shape == (100, 100, 100)
+    # Test against volume of ellipsoid +- center and subpixel issues
+    assert mask.sum() > 24 * 19 * 14 * 4/3 * np.pi
+    assert mask.sum() < 26 * 21 * 16 * 4/3 * np.pi
+
+def test_sphere():
+    mask = sphere(100, 50, 0, 1,0)
+    assert mask.shape == (100, 100, 100)
+    assert mask.sum() > 4/3 * np.pi * 24**3 # Volume of sphere
+    assert mask.sum() < 4/3 * np.pi * 26**3 # Volume of sphere
+
+def test_tube():
+    mask = tube(100, 50, 50, 0, 0, 1)
+    assert mask.shape == (100, 100, 100)
+    assert mask.sum() > np.pi * 24**2 * 48 # Volume of tube
+    assert mask.sum() < np.pi * 26**2 * 52 # Volume of tube
\ No newline at end of file
diff --git a/tests/test_ttmask.py b/tests/test_ttmask.py
deleted file mode 100644
index 363b3e2..0000000
--- a/tests/test_ttmask.py
+++ /dev/null
@@ -1,2 +0,0 @@
-def test_something():
-    pass