Repository: osbuild/osbuild

Issue:
[Azure] Provision VM failed on Azure because 91-azure_datasource.cfg format is incorrect
**Describe the bug**
In Image Builder images, the datasource_list in /etc/cloud/cloud.cfg.d/91-azure_datasource.cfg is:
```
...
datasource_list:
 - Azure
```
which cannot be parsed by cloud-init.
```
2024-01-18 10:27:12,433 - stages.py[DEBUG]: Using distro class <class 'cloudinit.distros.rhel.Distro'>
2024-01-18 10:27:12,441 - __init__.py[DEBUG]: Looking for data source in: [{'datasource_list': None}, 'None'], via packages ['', 'cloudinit.sources'] that matches dependencies ['FILESYSTEM', 'NETWORK']
2024-01-18 10:27:12,441 - util.py[WARNING]: failed stage init
2024-01-18 10:27:12,441 - util.py[DEBUG]: failed stage init
Traceback (most recent call last):
  File "/usr/lib/python3.6/site-packages/cloudinit/cmd/main.py", line 781, in status_wrapper
    ret = functor(name, args)
  File "/usr/lib/python3.6/site-packages/cloudinit/cmd/main.py", line 394, in main_init
    init.fetch(existing=existing)
  File "/usr/lib/python3.6/site-packages/cloudinit/stages.py", line 493, in fetch
    return self._get_data_source(existing=existing)
  File "/usr/lib/python3.6/site-packages/cloudinit/stages.py", line 367, in _get_data_source
    self.reporter,
  File "/usr/lib/python3.6/site-packages/cloudinit/sources/__init__.py", line 1001, in find_source
    ds_list = list_sources(cfg_list, ds_deps, pkg_list)
  File "/usr/lib/python3.6/site-packages/cloudinit/sources/__init__.py", line 1047, in list_sources
    ds_name = importer.match_case_insensitive_module_name(ds)
  File "/usr/lib/python3.6/site-packages/cloudinit/importer.py", line 40, in match_case_insensitive_module_name
    if "nocloud-net" == mod_name.lower():
AttributeError: 'dict' object has no attribute 'lower'
```
**Environment**
 * RHEL-7.9/8.10/RHEL-9.4
 * cloud-init-23.4-2.el8/el9
* (This issue doesn't exist in RHEL-7.9 itself, but if leapp upgrade 7.9 to RHEL-8.10, the wrong datasource_list will make cloud-init fail in RHEL-8.10. So suggest to also correct it in RHEL-7.9 images.)

**To Reproduce**
Steps to reproduce the behavior:
1. Create a VM from the latest RHEL-8.10/9.4 vhd images in Azure. 
The VM provision failed because cloud-init cannot parse the datasource_list and raise an exception.

**Expected behavior**
cloud-init service can start successfully.

**Additional context**
The /etc/cloud/cloud.cfg.d/91-azure_datasource.cfg in the Azure Marketplace RHEL images are correct:
```
datasource_list: [ Azure ]
datasource:
   Azure:
      apply_network_config: False
```


Discussion:

Patch:
diff --git a/stages/org.osbuild.cloud-init b/stages/org.osbuild.cloud-init
index 6361691c..aa986948 100755
--- a/stages/org.osbuild.cloud-init
+++ b/stages/org.osbuild.cloud-init
@@ -142,12 +142,34 @@ SCHEMA = r"""
 """
 
 
+# Class representing the 'datasource_list' configuration option,
+# allowing to define a custom YAML dumper for it.
+class DatasourceList(list):
+    pass
+
+
+# Dedicated YAML dumper for the `DatasourceList` class.
+# Make sure that the `DatasourceList` always uses flow style.
+# https://cloudinit.readthedocs.io/en/latest/reference/base_config_reference.html#datasource-list
+# Specifically, it says for the `datasource_list` key:
+# "This key is unique in that it uses a subset of YAML syntax.
+#  It requires that the key and its contents, a list,
+#  must share a single line - no newlines."
+def datasourcelist_presenter(dumper, data):
+    return dumper.represent_sequence("tag:yaml.org,2002:seq", data, flow_style=True)
+
+
 # Writes the passed `config` object as is into the configuration file in YAML format.
 # The validity of the `config` content is assured by the SCHEMA.
 def main(tree, options):
     filename = options.get("filename")
     config = options.get("config", {})
 
+    datasource_list = config.get("datasource_list")
+    if datasource_list:
+        config["datasource_list"] = DatasourceList(datasource_list)
+    yaml.add_representer(DatasourceList, datasourcelist_presenter)
+
     config_files_dir = f"{tree}/etc/cloud/cloud.cfg.d"
 
     with open(f"{config_files_dir}/{filename}", "w", encoding="utf8") as f:


Test patch:
diff --git a/stages/test/conftest.py b/stages/test/conftest.py
new file mode 100644
index 00000000..f2f9d5a7
--- /dev/null
+++ b/stages/test/conftest.py
@@ -0,0 +1,45 @@
+import os
+import pathlib
+from types import ModuleType
+
+import pytest
+
+import osbuild.meta
+from osbuild.testutil.imports import import_module_from_path
+
+
+@pytest.fixture
+def stage_module(request: pytest.FixtureRequest) -> ModuleType:
+    """stage_module is a fixture that imports a stage module by its name
+    defined in STAGE_NAME in the test module.
+    """
+    if not hasattr(request.module, "STAGE_NAME"):
+        raise ValueError("stage_module fixture must be used in a test module that defines STAGE_NAME")
+
+    stage_name = request.module.STAGE_NAME
+    caller_dir = pathlib.Path(request.node.fspath).parent
+    module_path = caller_dir.parent / stage_name
+    return import_module_from_path("stage", os.fspath(module_path))
+
+
+@pytest.fixture
+def stage_schema(request: pytest.FixtureRequest) -> osbuild.meta.Schema:
+    """stage_schema is a fixture returns the schema for a stage module.
+
+    This fixture may be indirectly parametrized with the stage schema version.
+    If the schema version is not specified, the version "2" is assumed.
+    The stage name must be defined in STAGE_NAME in the test module.
+    """
+    if hasattr(request, "param") and not isinstance(request.param, str):
+        raise ValueError(
+            "stage_schema fixture may be indirectly parametrized only with the stage schema version string")
+
+    if not hasattr(request.module, "STAGE_NAME"):
+        raise ValueError("stage_schema fixture must be used in a test module that defines STAGE_NAME")
+
+    stage_name = request.module.STAGE_NAME
+    schema_version = request.param if hasattr(request, "param") else "2"
+    caller_dir = pathlib.Path(request.node.fspath).parent
+    root = caller_dir.parent.parent
+    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", stage_name)
+    return osbuild.meta.Schema(mod_info.get_schema(version=schema_version), stage_name)
diff --git a/stages/test/test_autotailor.py b/stages/test/test_autotailor.py
index a42d0282..3fdd07c0 100644
--- a/stages/test/test_autotailor.py
+++ b/stages/test/test_autotailor.py
@@ -1,14 +1,11 @@
 #!/usr/bin/python3
 
-import os.path
 import sys
 from unittest.mock import call, patch
 
 import pytest
 
-import osbuild.meta
 from osbuild import testutil
-from osbuild.testutil.imports import import_module_from_path
 
 TEST_INPUT = [
     (
@@ -36,10 +33,13 @@ TEST_INPUT = [
 ]
 
 
+STAGE_NAME = "org.osbuild.oscap.autotailor"
+
+
 @pytest.fixture(name="fake_input")
 def fake_input_fixture():
     return {
-        "name": "org.osbuild.oscap.autotailor",
+        "name": STAGE_NAME,
         "options": {
             "filepath": "/some/filepath",
             "config": {
@@ -51,26 +51,12 @@ def fake_input_fixture():
     }
 
 
-def schema_validate_stage_oscap_autotailor(fake_input, test_data):
-    name = "org.osbuild.oscap.autotailor"
-    version = "1"
-    root = os.path.join(os.path.dirname(__file__), "../..")
-    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", name)
-    schema = osbuild.meta.Schema(mod_info.get_schema(version=version), name)
-    test_input = fake_input
-    test_input["options"]["config"].update(test_data)
-    return schema.validate(test_input)
-
-
 @pytest.mark.parametrize("test_overrides,expected", TEST_INPUT)
 @patch("subprocess.run")
-def test_oscap_autotailor_overrides_smoke(mock_subprocess_run, fake_input, test_overrides, expected):
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.oscap.autotailor")
-    stage = import_module_from_path("stage", stage_path)
-
+def test_oscap_autotailor_overrides_smoke(mock_subprocess_run, fake_input, stage_module, test_overrides, expected):
     options = fake_input["options"]
     options["config"]["overrides"] = test_overrides
-    stage.main("/some/sysroot", options)
+    stage_module.main("/some/sysroot", options)
 
     assert mock_subprocess_run.call_args_list == [
         call(["/usr/bin/autotailor", "--output", "/some/sysroot/some/filepath",
@@ -109,8 +95,10 @@ def test_oscap_autotailor_overrides_smoke(mock_subprocess_run, fake_input, test_
         }, " is not of type 'string', 'integer'"),
     ],
 )
-def test_schema_validation_oscap_autotailor(fake_input, test_data, expected_err):
-    res = schema_validate_stage_oscap_autotailor(fake_input, test_data)
+@pytest.mark.parametrize("stage_schema", ["1"], indirect=True)
+def test_schema_validation_oscap_autotailor(fake_input, stage_schema, test_data, expected_err):
+    fake_input["options"]["config"].update(test_data)
+    res = stage_schema.validate(fake_input)
 
     assert res.valid is False
     testutil.assert_jsonschema_error_contains(res, expected_err, expected_num_errs=1)
diff --git a/stages/test/test_bootupd.py b/stages/test/test_bootupd.py
index f9d3bfd3..06cc485b 100644
--- a/stages/test/test_bootupd.py
+++ b/stages/test/test_bootupd.py
@@ -1,13 +1,12 @@
 #!/usr/bin/python3
 
-import os.path
 from unittest.mock import call, patch
 
 import pytest
 
-import osbuild.meta
 from osbuild import testutil
-from osbuild.testutil.imports import import_module_from_path
+
+STAGE_NAME = "org.osbuild.bootupd"
 
 
 @pytest.mark.parametrize("test_data,expected_err", [
@@ -37,19 +36,14 @@ from osbuild.testutil.imports import import_module_from_path
     }, "")
 
 ])
-def test_bootupd_schema_validation(test_data, expected_err):
-    name = "org.osbuild.bootupd"
-    root = os.path.join(os.path.dirname(__file__), "../..")
-    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", name)
-    schema = osbuild.meta.Schema(mod_info.get_schema(version="2"), name)
-
+def test_bootupd_schema_validation(stage_schema, test_data, expected_err):
     test_input = {
-        "type": "org.osbuild.bootupd",
+        "type": STAGE_NAME,
         "options": {
         }
     }
     test_input["options"].update(test_data)
-    res = schema.validate(test_input)
+    res = stage_schema.validate(test_input)
 
     if expected_err == "":
         assert res.valid is True, f"err: {[e.as_dict() for e in res.errors]}"
@@ -64,10 +58,7 @@ def test_bootupd_schema_validation(test_data, expected_err):
     ({"static-configs": True}, ["--with-static-configs"]),
 ])
 @patch("subprocess.run")
-def test_bootupd_defaults(mocked_run, test_options, expected_bootupd_opts):
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.bootupd")
-    stage = import_module_from_path("bootupd_stage", stage_path)
-
+def test_bootupd_defaults(mocked_run, stage_module, test_options, expected_bootupd_opts):
     args = {
         "paths": {
             "mounts": "/run/osbuild/mounts",
@@ -79,8 +70,8 @@ def test_bootupd_defaults(mocked_run, test_options, expected_bootupd_opts):
         },
     }
     options = test_options
-    with patch.object(stage, "bind_mounts") as mocked_bind_mounts:
-        stage.main(args, options)
+    with patch.object(stage_module, "bind_mounts") as mocked_bind_mounts:
+        stage_module.main(args, options)
 
     assert len(mocked_bind_mounts.call_args_list) == 1
     assert mocked_run.call_args_list == [call(["chroot",
@@ -93,12 +84,9 @@ def test_bootupd_defaults(mocked_run, test_options, expected_bootupd_opts):
 
 
 @patch("subprocess.run")
-def test_bootupd_bind_mounts(mocked_run):
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.bootupd")
-    stage = import_module_from_path("bootupd_stage", stage_path)
-
+def test_bootupd_bind_mounts(mocked_run, stage_module):
     dst = "/run/osbuild/mounts/"
-    with stage.bind_mounts(['/dev', '/proc', '/sys', '/run', '/var', '/tmp'], dst):
+    with stage_module.bind_mounts(['/dev', '/proc', '/sys', '/run', '/var', '/tmp'], dst):
         pass
     assert mocked_run.call_args_list == [
         call(["mount", "--rbind", "/dev", "/run/osbuild/mounts/dev"], check=True),
diff --git a/stages/test/test_cloud-init.py b/stages/test/test_cloud-init.py
new file mode 100644
index 00000000..9a7f8bc9
--- /dev/null
+++ b/stages/test/test_cloud-init.py
@@ -0,0 +1,26 @@
+#!/usr/bin/python3
+
+import os.path
+
+STAGE_NAME = "org.osbuild.cloud-init"
+
+
+# Test that the 'datasource_list' option is dumped as a single line
+# in the configuration file.
+def test_cloud_init_datasource_list_single_line(tmp_path, stage_module):
+    treedir = tmp_path / "tree"
+    confpath = treedir / "etc/cloud/cloud.cfg.d/datasource_list.cfg"
+    confpath.parent.mkdir(parents=True, exist_ok=True)
+
+    options = {
+        "filename": confpath.name,
+        "config": {
+            "datasource_list": [
+                "Azure",
+            ]
+        }
+    }
+
+    stage_module.main(treedir, options)
+    assert os.path.exists(confpath)
+    assert confpath.read_text() == "datasource_list: [Azure]\n"
diff --git a/stages/test/test_container_deploy.py b/stages/test/test_container_deploy.py
index b76c919e..285e82a7 100644
--- a/stages/test/test_container_deploy.py
+++ b/stages/test/test_container_deploy.py
@@ -10,7 +10,8 @@ from unittest.mock import call, patch
 import pytest
 
 from osbuild.testutil import has_executable, make_fake_tree
-from osbuild.testutil.imports import import_module_from_path
+
+STAGE_NAME = "org.osbuild.container-deploy"
 
 
 def make_container(tmp_path, tag, fake_content, base="scratch"):
@@ -32,10 +33,7 @@ def make_container(tmp_path, tag, fake_content, base="scratch"):
 
 @pytest.mark.skipif(os.getuid() != 0, reason="needs root")
 @pytest.mark.skipif(not has_executable("podman"), reason="no podman executable")
-def test_container_deploy_integration(tmp_path):
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.container-deploy")
-    stage = import_module_from_path("stage", stage_path)
-
+def test_container_deploy_integration(tmp_path, stage_module):
     # build two containers and overlay files to test for
     # https://github.com/containers/storage/issues/1779
     base_tag = "cont-base-" + "".join(random.choices(string.digits, k=12))
@@ -71,7 +69,7 @@ def test_container_deploy_integration(tmp_path):
     options = {}
 
     with patch("os.makedirs", wraps=os.makedirs) as mocked_makedirs:
-        stage.main(inputs, output_dir, options)
+        stage_module.main(inputs, output_dir, options)
 
     assert output_dir.exists()
     assert (output_dir / "file1").read_bytes() == b"file1 from final layer"
@@ -81,10 +79,7 @@ def test_container_deploy_integration(tmp_path):
 
 @pytest.mark.skipif(os.getuid() != 0, reason="needs root")
 @pytest.mark.skipif(not has_executable("podman"), reason="no podman executable")
-def test_container_deploy_exclude(tmp_path):
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.container-deploy")
-    stage = import_module_from_path("stage", stage_path)
-
+def test_container_deploy_exclude(tmp_path, stage_module):
     base_tag = "cont-base-" + "".join(random.choices(string.digits, k=12))
     make_container(tmp_path, base_tag, {
         "file1": "file1 content",
@@ -125,7 +120,7 @@ def test_container_deploy_exclude(tmp_path):
     }
     output_dir = tmp_path / "output"
 
-    stage.main(inputs, output_dir, options)
+    stage_module.main(inputs, output_dir, options)
     assert output_dir.exists()
     assert (output_dir / "file1").read_bytes() == b"file1 content"
     assert not (output_dir / "file2").exists()
diff --git a/stages/test/test_containers_storage_conf.py b/stages/test/test_containers_storage_conf.py
index 36f322c6..2aed0a5e 100644
--- a/stages/test/test_containers_storage_conf.py
+++ b/stages/test/test_containers_storage_conf.py
@@ -10,9 +10,7 @@ try:
 except ModuleNotFoundError:
     import pytoml as toml
 
-import osbuild.meta
 from osbuild import testutil
-from osbuild.testutil.imports import import_module_from_path
 
 TEST_INPUT = [
     ({}, {"additionalimagestores": ["/path/to/store"]}, [("storage.options.additionalimagestores", ["/path/to/store"])]),
@@ -21,16 +19,17 @@ TEST_INPUT = [
 ]
 
 
+STAGE_NAME = "org.osbuild.containers.storage.conf"
+
+
 @pytest.mark.parametrize("test_filename", ["/etc/containers/storage.conf", "/usr/share/containers/storage.conf"])
 @pytest.mark.parametrize("test_storage,test_options,expected", TEST_INPUT)
-def test_containers_storage_conf_integration(tmp_path, test_filename, test_storage, test_options, expected):  # pylint: disable=unused-argument
+def test_containers_storage_conf_integration(tmp_path, stage_module, test_filename, test_storage, test_options, expected):  # pylint: disable=unused-argument
     treedir = os.path.join(tmp_path, "tree")
     confdir, confname = os.path.split(test_filename.lstrip("/"))
     confdir = os.path.join(treedir, confdir)
     os.makedirs(confdir, exist_ok=True)
 
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.containers.storage.conf")
-    stage = import_module_from_path("stage", stage_path)
     options = {
         "filename": test_filename,
         "config": {
@@ -44,7 +43,7 @@ def test_containers_storage_conf_integration(tmp_path, test_filename, test_stora
     options["config"]["storage"].update(test_storage)
     options["config"]["storage"]["options"].update(test_options)
 
-    stage.main(treedir, options)
+    stage_module.main(treedir, options)
 
     confpath = os.path.join(confdir, confname)
     assert os.path.exists(confpath)
@@ -95,15 +94,9 @@ def test_containers_storage_conf_integration(tmp_path, test_filename, test_stora
         ({}, {"options": {"pull_options": {"use_hard_links": True}}}, "is not of type 'string'"),
     ],
 )
-def test_schema_validation_containers_storage_conf(test_data, storage_test_data, expected_err):
-    name = "org.osbuild.containers.storage.conf"
-    version = "2"
-    root = os.path.join(os.path.dirname(__file__), "../..")
-    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", name)
-    schema = osbuild.meta.Schema(mod_info.get_schema(version=version), name)
-
+def test_schema_validation_containers_storage_conf(stage_schema, test_data, storage_test_data, expected_err):
     test_input = {
-        "type": "org.osbuild.containers.storage.conf",
+        "type": STAGE_NAME,
         "options": {
             "config": {
                 "storage": {
@@ -114,7 +107,7 @@ def test_schema_validation_containers_storage_conf(test_data, storage_test_data,
 
     test_input["options"].update(test_data)
     test_input["options"]["config"]["storage"].update(storage_test_data)
-    res = schema.validate(test_input)
+    res = stage_schema.validate(test_input)
 
     if expected_err == "":
         assert res.valid is True, f"err: {[e.as_dict() for e in res.errors]}"
diff --git a/stages/test/test_erofs.py b/stages/test/test_erofs.py
index 45ddfd2f..3b41519f 100644
--- a/stages/test/test_erofs.py
+++ b/stages/test/test_erofs.py
@@ -6,10 +6,8 @@ from unittest import mock
 
 import pytest
 
-import osbuild.meta
 from osbuild import testutil
 from osbuild.testutil import has_executable, make_fake_input_tree
-from osbuild.testutil.imports import import_module_from_path
 
 TEST_INPUT = [
     ({}, []),
@@ -20,9 +18,12 @@ TEST_INPUT = [
 ]
 
 
+STAGE_NAME = "org.osbuild.erofs"
+
+
 @pytest.mark.skipif(not has_executable("mkfs.erofs"), reason="no mkfs.erofs")
 @pytest.mark.parametrize("test_options,expected", TEST_INPUT)
-def test_erofs_integration(tmp_path, test_options, expected):  # pylint: disable=unused-argument
+def test_erofs_integration(tmp_path, stage_module, test_options, expected):  # pylint: disable=unused-argument
     fake_input_tree = make_fake_input_tree(tmp_path, {
         "/file-in-root.txt": "other content",
         "/subdir/file-in-subdir.txt": "subdir content",
@@ -32,15 +33,13 @@ def test_erofs_integration(tmp_path, test_options, expected):  # pylint: disable
             "path": fake_input_tree,
         }
     }
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.erofs")
-    stage = import_module_from_path("erofs_stage", stage_path)
     filename = "some-file.img"
     options = {
         "filename": filename,
     }
     options.update(test_options)
 
-    stage.main(inputs, tmp_path, options)
+    stage_module.main(inputs, tmp_path, options)
 
     img_path = os.path.join(tmp_path, "some-file.img")
     assert os.path.exists(img_path)
@@ -56,7 +55,7 @@ def test_erofs_integration(tmp_path, test_options, expected):  # pylint: disable
 
 @mock.patch("subprocess.run")
 @pytest.mark.parametrize("test_options,expected", TEST_INPUT)
-def test_erofs(mock_run, tmp_path, test_options, expected):
+def test_erofs(mock_run, tmp_path, stage_module, test_options, expected):
     fake_input_tree = make_fake_input_tree(tmp_path, {
         "/some-dir/some-file.txt": "content",
     })
@@ -65,15 +64,13 @@ def test_erofs(mock_run, tmp_path, test_options, expected):
             "path": fake_input_tree,
         }
     }
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.erofs")
-    stage = import_module_from_path("erofs_stage", stage_path)
     filename = "some-file.img"
     options = {
         "filename": filename,
     }
     options.update(test_options)
 
-    stage.main(inputs, tmp_path, options)
+    stage_module.main(inputs, tmp_path, options)
 
     expected = [
         "mkfs.erofs",
@@ -91,20 +88,15 @@ def test_erofs(mock_run, tmp_path, test_options, expected):
     # good
     ({"compression": {"method": "lz4"}}, ""),
 ])
-def test_schema_validation_erofs(test_data, expected_err):
-    name = "org.osbuild.erofs"
-    root = os.path.join(os.path.dirname(__file__), "../..")
-    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", name)
-    schema = osbuild.meta.Schema(mod_info.get_schema(version="2"), name)
-
+def test_schema_validation_erofs(stage_schema, test_data, expected_err):
     test_input = {
-        "type": "org.osbuild.erofs",
+        "type": STAGE_NAME,
         "options": {
             "filename": "some-filename.img",
         }
     }
     test_input["options"].update(test_data)
-    res = schema.validate(test_input)
+    res = stage_schema.validate(test_input)
 
     if expected_err == "":
         assert res.valid is True, f"err: {[e.as_dict() for e in res.errors]}"
diff --git a/stages/test/test_grub2.py b/stages/test/test_grub2.py
index b95cbd46..f3197baf 100644
--- a/stages/test/test_grub2.py
+++ b/stages/test/test_grub2.py
@@ -1,15 +1,11 @@
 #!/usr/bin/python3
 
-import os.path
-
 from osbuild.testutil import make_fake_tree
-from osbuild.testutil.imports import import_module_from_path
 
+STAGE_NAME = "org.osbuild.grub2"
 
-def test_grub2_copy_efi_data(tmp_path):
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.grub2")
-    stage = import_module_from_path("stage", stage_path)
 
+def test_grub2_copy_efi_data(tmp_path, stage_module):
     fake_efi_src_dir = tmp_path / "fake-efi/EFI"
     make_fake_tree(fake_efi_src_dir, {
         "fedora/a.shim": "fake shim",
@@ -27,6 +23,6 @@ def test_grub2_copy_efi_data(tmp_path):
         },
     }
     fake_tree = tmp_path / "tree"
-    stage.main(fake_tree, test_options)
+    stage_module.main(fake_tree, test_options)
     assert (fake_tree / "boot/efi/EFI/fedora/a.shim").exists()
     assert (fake_tree / "boot/efi/EFI/BOOT/BOOTX64.EFI").exists()
diff --git a/stages/test/test_kickstart.py b/stages/test/test_kickstart.py
index 4a7957e9..d4aab3a8 100644
--- a/stages/test/test_kickstart.py
+++ b/stages/test/test_kickstart.py
@@ -5,10 +5,8 @@ import subprocess
 
 import pytest
 
-import osbuild.meta
 from osbuild import testutil
 from osbuild.testutil import has_executable
-from osbuild.testutil.imports import import_module_from_path
 
 TEST_INPUT = [
     ({"lang": "en_US.UTF-8"}, "lang en_US.UTF-8"),
@@ -216,39 +214,31 @@ TEST_INPUT = [
 ]
 
 
-def schema_validate_kickstart_stage(test_data):
-    name = "org.osbuild.kickstart"
-    version = "1"
-    root = os.path.join(os.path.dirname(__file__), "../..")
-    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", name)
-    schema = osbuild.meta.Schema(mod_info.get_schema(version=version), name)
-    test_input = {
-        "name": "org.osbuild.kickstart",
-        "options": {
-            "path": "some-path",
-        }
-    }
-    test_input["options"].update(test_data)
-    return schema.validate(test_input)
+STAGE_NAME = "org.osbuild.kickstart"
 
 
+@pytest.mark.parametrize("stage_schema", ["1"], indirect=True)
 @pytest.mark.parametrize("test_input,expected", TEST_INPUT)
-def test_kickstart_test_cases_valid(test_input, expected):  # pylint: disable=unused-argument
+def test_kickstart_test_cases_valid(stage_schema, test_input, expected):  # pylint: disable=unused-argument
     """ ensure all test inputs are valid """
-    res = schema_validate_kickstart_stage(test_input)
+    test_data = {
+        "name": STAGE_NAME,
+        "options": {
+            "path": "some-path",
+        }
+    }
+    test_data["options"].update(test_input)
+    res = stage_schema.validate(test_data)
     assert res.valid is True, f"input: {test_input}\nerr: {[e.as_dict() for e in res.errors]}"
 
 
 @pytest.mark.parametrize("test_input,expected", TEST_INPUT)
-def test_kickstart_write(tmp_path, test_input, expected):
-    ks_stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.kickstart")
-    ks_stage = import_module_from_path("ks_stage", ks_stage_path)
-
+def test_kickstart_write(tmp_path, stage_module, test_input, expected):
     ks_path = "kickstart/kfs.cfg"
     options = {"path": ks_path}
     options.update(test_input)
 
-    ks_stage.main(tmp_path, options)
+    stage_module.main(tmp_path, options)
 
     ks_path = os.path.join(tmp_path, ks_path)
     with open(ks_path, encoding="utf-8") as fp:
@@ -258,15 +248,12 @@ def test_kickstart_write(tmp_path, test_input, expected):
 
 @pytest.mark.skipif(not has_executable("ksvalidator"), reason="`ksvalidator` is required")
 @pytest.mark.parametrize("test_input,expected", TEST_INPUT)
-def test_kickstart_valid(tmp_path, test_input, expected):  # pylint: disable=unused-argument
-    ks_stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.kickstart")
-    ks_stage = import_module_from_path("ks_stage", ks_stage_path)
-
+def test_kickstart_valid(tmp_path, stage_module, test_input, expected):  # pylint: disable=unused-argument
     ks_path = "kickstart/kfs.cfg"
     options = {"path": ks_path}
     options.update(test_input)
 
-    ks_stage.main(tmp_path, options)
+    stage_module.main(tmp_path, options)
 
     ks_path = os.path.join(tmp_path, ks_path)
 
@@ -275,7 +262,7 @@ def test_kickstart_valid(tmp_path, test_input, expected):  # pylint: disable=unu
 
 
 @pytest.mark.parametrize(
-    "test_data,expected_err",
+    "test_input,expected_err",
     [
         # BAD pattern, ensure some obvious ways to write arbitrary
         # kickstart files will not work
@@ -362,8 +349,16 @@ def test_kickstart_valid(tmp_path, test_input, expected):  # pylint: disable=unu
 
     ],
 )
-def test_schema_validation_bad_apples(test_data, expected_err):
-    res = schema_validate_kickstart_stage(test_data)
+@pytest.mark.parametrize("stage_schema", ["1"], indirect=True)
+def test_schema_validation_bad_apples(stage_schema, test_input, expected_err):
+    test_data = {
+        "name": STAGE_NAME,
+        "options": {
+            "path": "some-path",
+        }
+    }
+    test_data["options"].update(test_input)
+    res = stage_schema.validate(test_data)
 
     assert res.valid is False
     testutil.assert_jsonschema_error_contains(res, expected_err, expected_num_errs=1)
diff --git a/stages/test/test_machine-id.py b/stages/test/test_machine-id.py
index b5ad2bda..741cfb49 100644
--- a/stages/test/test_machine-id.py
+++ b/stages/test/test_machine-id.py
@@ -1,20 +1,13 @@
 #!/usr/bin/python3
 
 import os
-import pathlib
 import unittest.mock
 
 import pytest
 
-import osbuild.meta
 from osbuild import testutil
-from osbuild.testutil.imports import import_module_from_path
 
-
-def stage(stage_name):
-    test_dir = pathlib.Path(__file__).parent
-    stage_path = pathlib.Path(test_dir) / f"../org.osbuild.{stage_name}"
-    return import_module_from_path("stage", os.fspath(stage_path))
+STAGE_NAME = "org.osbuild.machine-id"
 
 
 @pytest.fixture(name="machine_id_path")
@@ -25,30 +18,35 @@ def machine_id_path_fixture(tmp_path):
 
 
 @pytest.mark.parametrize("already_has_etc_machine_id", [True, False])
-def test_machine_id_first_boot_yes(tmp_path, machine_id_path, already_has_etc_machine_id):
+def test_machine_id_first_boot_yes(tmp_path, stage_module, machine_id_path, already_has_etc_machine_id):
     if already_has_etc_machine_id:
         machine_id_path.touch()
 
-    stage("machine-id").main(tmp_path, {"first-boot": "yes"})
+    stage_module.main(tmp_path, {"first-boot": "yes"})
     assert machine_id_path.read_bytes() == b"uninitialized\n"
 
 
 @pytest.mark.parametrize("already_has_etc_machine_id", [True, False])
-def test_machine_id_first_boot_no(tmp_path, machine_id_path, already_has_etc_machine_id):
+def test_machine_id_first_boot_no(tmp_path, stage_module, machine_id_path, already_has_etc_machine_id):
     if already_has_etc_machine_id:
         machine_id_path.write_bytes(b"\x01\x02\x03")
 
-    stage("machine-id").main(tmp_path, {"first-boot": "no"})
+    stage_module.main(tmp_path, {"first-boot": "no"})
     assert machine_id_path.stat().st_size == 0
 
 
 @pytest.mark.parametrize("already_has_etc_machine_id", [True, False])
 @unittest.mock.patch("builtins.print")
-def test_machine_id_first_boot_preserve(mock_print, tmp_path, machine_id_path, already_has_etc_machine_id):
+def test_machine_id_first_boot_preserve(
+        mock_print,
+        tmp_path,
+        stage_module,
+        machine_id_path,
+        already_has_etc_machine_id):
     if already_has_etc_machine_id:
         machine_id_path.write_bytes(b"\x01\x02\x03")
 
-    ret = stage("machine-id").main(tmp_path, {"first-boot": "preserve"})
+    ret = stage_module.main(tmp_path, {"first-boot": "preserve"})
     if already_has_etc_machine_id:
         assert os.stat(machine_id_path).st_size == 3
     else:
@@ -59,23 +57,19 @@ def test_machine_id_first_boot_preserve(mock_print, tmp_path, machine_id_path, a
 @pytest.mark.parametrize("test_data,expected_err", [
     ({"first-boot": "invalid-option"}, "'invalid-option' is not one of "),
 ])
-def test_machine_id_schema_validation(test_data, expected_err):
-    name = "org.osbuild.machine-id"
-    root = pathlib.Path(__file__).parents[2]
-    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", name)
-    schema = osbuild.meta.Schema(mod_info.get_schema(), name)
-
+@pytest.mark.parametrize("stage_schema", ["1"], indirect=True)
+def test_machine_id_schema_validation(stage_schema, test_data, expected_err):
     test_input = {
-        "name": "org.osbuild.machine-id",
+        "name": STAGE_NAME,
         "options": {},
     }
     test_input["options"].update(test_data)
-    res = schema.validate(test_input)
+    res = stage_schema.validate(test_input)
 
     assert res.valid is False
     testutil.assert_jsonschema_error_contains(res, expected_err, expected_num_errs=1)
 
 
-def test_machine_id_first_boot_unknown(tmp_path):
+def test_machine_id_first_boot_unknown(tmp_path, stage_module):
     with pytest.raises(ValueError, match=r"unexpected machine-id mode 'invalid'"):
-        stage("machine-id").main(tmp_path, {"first-boot": "invalid"})
+        stage_module.main(tmp_path, {"first-boot": "invalid"})
diff --git a/stages/test/test_mkfs_ext4.py b/stages/test/test_mkfs_ext4.py
index bb3433e4..e83666da 100644
--- a/stages/test/test_mkfs_ext4.py
+++ b/stages/test/test_mkfs_ext4.py
@@ -7,10 +7,10 @@ from unittest import mock
 
 import pytest
 
-import osbuild.meta
 from osbuild import testutil
 from osbuild.testutil import has_executable
-from osbuild.testutil.imports import import_module_from_path
+
+STAGE_NAME = "org.osbuild.mkfs.ext4"
 
 
 @pytest.mark.parametrize("test_data,expected_err", [
@@ -26,14 +26,9 @@ from osbuild.testutil.imports import import_module_from_path
     # uuid but our schema is not strict enough right now
     ({"uuid": "some-uuid"}, ""),
 ])
-def test_schema_validation_mkfs_ext4(test_data, expected_err):
-    name = "org.osbuild.mkfs.ext4"
-    root = os.path.join(os.path.dirname(__file__), "../..")
-    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", name)
-    schema = osbuild.meta.Schema(mod_info.get_schema(version="2"), name)
-
+def test_schema_validation_mkfs_ext4(stage_schema, test_data, expected_err):
     test_input = {
-        "type": "org.osbuild.mkfs.ext4",
+        "type": STAGE_NAME,
         "devices": {
             "device": {
                 "path": "some-path",
@@ -43,7 +38,7 @@ def test_schema_validation_mkfs_ext4(test_data, expected_err):
         }
     }
     test_input["options"].update(test_data)
-    res = schema.validate(test_input)
+    res = stage_schema.validate(test_input)
 
     if expected_err == "":
         assert res.valid is True, f"err: {[e.as_dict() for e in res.errors]}"
@@ -53,10 +48,7 @@ def test_schema_validation_mkfs_ext4(test_data, expected_err):
 
 
 @pytest.mark.skipif(not has_executable("mkfs.ext4"), reason="need mkfs.ext4")
-def test_mkfs_ext4_integration(tmp_path):
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.mkfs.ext4")
-    stage = import_module_from_path("mkfs_ext4_stage", stage_path)
-
+def test_mkfs_ext4_integration(tmp_path, stage_module):
     fake_disk_path = tmp_path / "fake.img"
     with fake_disk_path.open("w"):
         os.truncate(os.fspath(fake_disk_path), 10 * 1024 * 1024)
@@ -69,7 +61,7 @@ def test_mkfs_ext4_integration(tmp_path):
     options = {
         "uuid": fake_uuid,
     }
-    stage.main(devices, options)
+    stage_module.main(devices, options)
     assert os.path.exists(fake_disk_path)
     output = subprocess.check_output([
         "dumpe2fs", fake_disk_path], encoding="utf-8")
@@ -88,10 +80,7 @@ def test_mkfs_ext4_integration(tmp_path):
     ({"verity": True, "orphan_file": True, "metadata_csum_seed": True}, ["-O", "verity", "-O", "orphan_file", "-O", "metadata_csum_seed"]),
 ])
 @mock.patch("subprocess.run")
-def test_mkfs_ext4_cmdline(mock_run, test_input, expected):
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.mkfs.ext4")
-    stage = import_module_from_path("mkfs_ext4_stage", stage_path)
-
+def test_mkfs_ext4_cmdline(mock_run, stage_module, test_input, expected):
     fake_disk_path = "/dev/xxd1"
     devices = {
         "device": {
@@ -103,7 +92,7 @@ def test_mkfs_ext4_cmdline(mock_run, test_input, expected):
         "uuid": fake_uuid,
     }
     options.update(test_input)
-    stage.main(devices, options)
+    stage_module.main(devices, options)
 
     expected = [
         "mkfs.ext4",
diff --git a/stages/test/test_ostree_post_copy.py b/stages/test/test_ostree_post_copy.py
index 8e7a1fbc..8ad70607 100644
--- a/stages/test/test_ostree_post_copy.py
+++ b/stages/test/test_ostree_post_copy.py
@@ -1,43 +1,23 @@
 #!/usr/bin/python3
 
-import os.path
 from unittest.mock import call, patch
 
 import pytest
 
-import osbuild.meta
 from osbuild import testutil
-from osbuild.testutil.imports import import_module_from_path
 
-
-def schema_validate_stage_ostree_post_copy(test_data):
-    name = "org.osbuild.ostree.post-copy"
-    version = "2"
-    root = os.path.join(os.path.dirname(__file__), "../..")
-    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", name)
-    schema = osbuild.meta.Schema(mod_info.get_schema(version=version), name)
-    test_input = {
-        "type": "org.osbuild.ostree.post-copy",
-        "options": {
-            "sysroot": "/some/sysroot",
-        },
-    }
-    test_input.update(test_data)
-    return schema.validate(test_input)
+STAGE_NAME = "org.osbuild.ostree.post-copy"
 
 
 @patch("osbuild.util.ostree.cli")
-def test_ostree_post_copy_smoke(mock_ostree_cli):
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.ostree.post-copy")
-    stage = import_module_from_path("stage", stage_path)
-
+def test_ostree_post_copy_smoke(mock_ostree_cli, stage_module):
     paths = {
         "mounts": "/run/osbuild/mounts",
     }
     options = {
         "sysroot": "/some/sysroot",
     }
-    stage.main(paths, options)
+    stage_module.main(paths, options)
 
     assert mock_ostree_cli.call_args_list == [
         call("admin", "post-copy", sysroot="/run/osbuild/mounts/some/sysroot")]
@@ -53,8 +33,15 @@ def test_ostree_post_copy_smoke(mock_ostree_cli):
         ({"mounts": "must-be-array"}, " is not of type 'array'"),
     ],
 )
-def test_schema_validation_ostree_post_copy(test_data, expected_err):
-    res = schema_validate_stage_ostree_post_copy(test_data)
+def test_schema_validation_ostree_post_copy(stage_schema, test_data, expected_err):
+    test_input = {
+        "type": STAGE_NAME,
+        "options": {
+            "sysroot": "/some/sysroot",
+        },
+    }
+    test_input.update(test_data)
+    res = stage_schema.validate(test_input)
 
     assert res.valid is False
     testutil.assert_jsonschema_error_contains(res, expected_err, expected_num_errs=1)
diff --git a/stages/test/test_selinux.py b/stages/test/test_selinux.py
index 81360ce0..469d46d4 100644
--- a/stages/test/test_selinux.py
+++ b/stages/test/test_selinux.py
@@ -5,26 +5,21 @@ from unittest.mock import call, patch
 
 import pytest
 
-import osbuild.meta
 from osbuild import testutil
-from osbuild.testutil.imports import import_module_from_path
 
+STAGE_NAME = "org.osbuild.selinux"
 
-def schema_validation_selinux(test_data, implicit_file_contexts=True):
-    name = "org.osbuild.selinux"
-    root = os.path.join(os.path.dirname(__file__), "../..")
-    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", name)
-    schema = osbuild.meta.Schema(mod_info.get_schema(version="1"), name)
 
+def get_test_input(test_data, implicit_file_contexts=True):
     test_input = {
-        "name": "org.osbuild.selinux",
+        "name": STAGE_NAME,
         "options": {}
     }
     if implicit_file_contexts:
         test_input["options"]["file_contexts"] = "some-context"
 
     test_input["options"].update(test_data)
-    return schema.validate(test_input)
+    return test_input
 
 
 @pytest.mark.parametrize("test_data,expected_err", [
@@ -36,8 +31,9 @@ def schema_validation_selinux(test_data, implicit_file_contexts=True):
     ({"labels": "xxx"}, "'xxx' is not of type 'object'"),
     ({"force_autorelabel": "foo"}, "'foo' is not of type 'boolean'"),
 ])
-def test_schema_validation_selinux(test_data, expected_err):
-    res = schema_validation_selinux(test_data)
+@pytest.mark.parametrize("stage_schema", ["1"], indirect=True)
+def test_schema_validation_selinux(stage_schema, test_data, expected_err):
+    res = stage_schema.validate(get_test_input(test_data))
     if expected_err == "":
         assert res.valid is True, f"err: {[e.as_dict() for e in res.errors]}"
     else:
@@ -45,23 +41,20 @@ def test_schema_validation_selinux(test_data, expected_err):
         testutil.assert_jsonschema_error_contains(res, expected_err, expected_num_errs=1)
 
 
-def test_schema_validation_selinux_file_context_required():
-    test_data = {}
-    res = schema_validation_selinux(test_data, implicit_file_contexts=False)
+@pytest.mark.parametrize("stage_schema", ["1"], indirect=True)
+def test_schema_validation_selinux_file_context_required(stage_schema):
+    res = stage_schema.validate(get_test_input({}, implicit_file_contexts=False))
     assert res.valid is False
     expected_err = "'file_contexts' is a required property"
     testutil.assert_jsonschema_error_contains(res, expected_err, expected_num_errs=1)
 
 
 @patch("osbuild.util.selinux.setfiles")
-def test_selinux_file_contexts(mocked_setfiles, tmp_path):
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.selinux")
-    stage = import_module_from_path("stage", stage_path)
-
+def test_selinux_file_contexts(mocked_setfiles, tmp_path, stage_module):
     options = {
         "file_contexts": "etc/selinux/thing",
     }
-    stage.main(tmp_path, options)
+    stage_module.main(tmp_path, options)
 
     assert len(mocked_setfiles.call_args_list) == 1
     assert mocked_setfiles.call_args_list == [
@@ -71,11 +64,8 @@ def test_selinux_file_contexts(mocked_setfiles, tmp_path):
 
 @patch("osbuild.util.selinux.setfilecon")
 @patch("osbuild.util.selinux.setfiles")
-def test_selinux_labels(mocked_setfiles, mocked_setfilecon, tmp_path):
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.selinux")
-    stage = import_module_from_path("stage", stage_path)
-
-    osbuild.testutil.make_fake_input_tree(tmp_path, {
+def test_selinux_labels(mocked_setfiles, mocked_setfilecon, tmp_path, stage_module):
+    testutil.make_fake_input_tree(tmp_path, {
         "/usr/bin/bootc": "I'm only an imposter",
     })
 
@@ -85,7 +75,7 @@ def test_selinux_labels(mocked_setfiles, mocked_setfilecon, tmp_path):
             "/tree/usr/bin/bootc": "system_u:object_r:install_exec_t:s0",
         }
     }
-    stage.main(tmp_path, options)
+    stage_module.main(tmp_path, options)
 
     assert len(mocked_setfiles.call_args_list) == 1
     assert len(mocked_setfilecon.call_args_list) == 1
@@ -95,15 +85,12 @@ def test_selinux_labels(mocked_setfiles, mocked_setfilecon, tmp_path):
 
 
 @patch("osbuild.util.selinux.setfiles")
-def test_selinux_force_autorelabel(mocked_setfiles, tmp_path):  # pylint: disable=unused-argument
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.selinux")
-    stage = import_module_from_path("stage", stage_path)
-
+def test_selinux_force_autorelabel(mocked_setfiles, tmp_path, stage_module):  # pylint: disable=unused-argument
     for enable_autorelabel in [False, True]:
         options = {
             "file_contexts": "etc/selinux/thing",
             "force_autorelabel": enable_autorelabel,
         }
-        stage.main(tmp_path, options)
+        stage_module.main(tmp_path, options)
 
         assert (tmp_path / ".autorelabel").exists() == enable_autorelabel
diff --git a/stages/test/test_skopeo.py b/stages/test/test_skopeo.py
index 907bba83..2b591e72 100644
--- a/stages/test/test_skopeo.py
+++ b/stages/test/test_skopeo.py
@@ -1,12 +1,11 @@
 #!/usr/bin/python3
 
-import os.path
-
 import pytest
 
-import osbuild.meta
 from osbuild import testutil
 
+STAGE_NAME = "org.osbuild.skopeo"
+
 
 @pytest.mark.parametrize("test_data,expected_err", [
     # bad
@@ -21,18 +20,13 @@ from osbuild import testutil
     # *inputs* and it'll be a no-op in the stage
     ({"destination": {"type": "containers-storage"}}, ""),
 ])
-def test_schema_validation_skopeo(test_data, expected_err):
-    name = "org.osbuild.skopeo"
-    root = os.path.join(os.path.dirname(__file__), "../..")
-    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", name)
-    schema = osbuild.meta.Schema(mod_info.get_schema(version="2"), name)
-
+def test_schema_validation_skopeo(stage_schema, test_data, expected_err):
     test_input = {
-        "type": "org.osbuild.skopeo",
+        "type": STAGE_NAME,
         "options": {},
     }
     test_input["options"].update(test_data)
-    res = schema.validate(test_input)
+    res = stage_schema.validate(test_input)
 
     if expected_err == "":
         assert res.valid is True, f"err: {[e.as_dict() for e in res.errors]}"
diff --git a/stages/test/test_systemd.py b/stages/test/test_systemd.py
index 92eef2e5..90f8e861 100644
--- a/stages/test/test_systemd.py
+++ b/stages/test/test_systemd.py
@@ -3,16 +3,13 @@
 import os
 import os.path
 
-from osbuild.testutil.imports import import_module_from_path
+STAGE_NAME = "org.osbuild.systemd"
 
 
-def test_systemd_masked_generators(tmp_path):
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.systemd")
-    stage = import_module_from_path("stage", stage_path)
-
+def test_systemd_masked_generators(tmp_path, stage_module):
     options = {"masked_generators": ["fake-generator"]}
 
-    stage.main(tmp_path, options)
+    stage_module.main(tmp_path, options)
 
     masked_generator_path = os.path.join(tmp_path, "etc/systemd/system-generators/fake-generator")
     assert os.path.lexists(masked_generator_path)
diff --git a/stages/test/test_xz.py b/stages/test/test_xz.py
index 79bbf74b..02415ce9 100644
--- a/stages/test/test_xz.py
+++ b/stages/test/test_xz.py
@@ -6,10 +6,10 @@ from unittest import mock
 
 import pytest
 
-import osbuild.meta
 from osbuild import testutil
 from osbuild.testutil import has_executable, make_fake_input_tree
-from osbuild.testutil.imports import import_module_from_path
+
+STAGE_NAME = "org.osbuild.xz"
 
 
 @pytest.mark.parametrize("test_data,expected_err", [
@@ -18,19 +18,14 @@ from osbuild.testutil.imports import import_module_from_path
     # good
     ({"filename": "image.xz"}, ""),
 ])
-def test_schema_validation_xz(test_data, expected_err):
-    name = "org.osbuild.xz"
-    root = os.path.join(os.path.dirname(__file__), "../..")
-    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", name)
-    schema = osbuild.meta.Schema(mod_info.get_schema(version="2"), name)
-
+def test_schema_validation_xz(stage_schema, test_data, expected_err):
     test_input = {
-        "type": "org.osbuild.xz",
+        "type": STAGE_NAME,
         "options": {
         }
     }
     test_input["options"].update(test_data)
-    res = schema.validate(test_input)
+    res = stage_schema.validate(test_input)
 
     if expected_err == "":
         assert res.valid is True, f"err: {[e.as_dict() for e in res.errors]}"
@@ -58,14 +53,12 @@ def fake_input(tmp_path):
 
 
 @pytest.mark.skipif(not has_executable("xz"), reason="no xz executable")
-def test_xz_integration(tmp_path, fake_input_tree):  # pylint: disable=unused-argument
+def test_xz_integration(tmp_path, stage_module, fake_input_tree):  # pylint: disable=unused-argument
     inputs = fake_input_tree[1]
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.xz")
-    stage = import_module_from_path("xz_stage", stage_path)
     options = {
         "filename": "image.txt.xz",
     }
-    stage.main(inputs, tmp_path, options)
+    stage_module.main(inputs, tmp_path, options)
 
     img_path = os.path.join(tmp_path, "image.txt.xz")
     assert os.path.exists(img_path)
@@ -75,16 +68,14 @@ def test_xz_integration(tmp_path, fake_input_tree):  # pylint: disable=unused-ar
 
 
 @mock.patch("subprocess.run")
-def test_xz_cmdline(mock_run, tmp_path, fake_input_tree):
+def test_xz_cmdline(mock_run, tmp_path, stage_module, fake_input_tree):
     fake_input_path = fake_input_tree[0]
     inputs = fake_input_tree[1]
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.xz")
-    stage = import_module_from_path("xz_stage", stage_path)
     filename = "image.txt.xz"
     options = {
         "filename": filename,
     }
-    stage.main(inputs, tmp_path, options)
+    stage_module.main(inputs, tmp_path, options)
 
     expected = [
         "xz",
diff --git a/stages/test/test_zstd.py b/stages/test/test_zstd.py
index ec702cd9..8cccde13 100644
--- a/stages/test/test_zstd.py
+++ b/stages/test/test_zstd.py
@@ -6,10 +6,10 @@ from unittest import mock
 
 import pytest
 
-import osbuild.meta
 from osbuild import testutil
 from osbuild.testutil import has_executable, make_fake_input_tree
-from osbuild.testutil.imports import import_module_from_path
+
+STAGE_NAME = "org.osbuild.zstd"
 
 
 @pytest.mark.parametrize("test_data,expected_err", [
@@ -18,19 +18,14 @@ from osbuild.testutil.imports import import_module_from_path
     # good
     ({"filename": "image.zst"}, ""),
 ])
-def test_schema_validation_zstd(test_data, expected_err):
-    name = "org.osbuild.zstd"
-    root = os.path.join(os.path.dirname(__file__), "../..")
-    mod_info = osbuild.meta.ModuleInfo.load(root, "Stage", name)
-    schema = osbuild.meta.Schema(mod_info.get_schema(version="2"), name)
-
+def test_schema_validation_zstd(stage_schema, test_data, expected_err):
     test_input = {
-        "type": "org.osbuild.zstd",
+        "type": STAGE_NAME,
         "options": {
         }
     }
     test_input["options"].update(test_data)
-    res = schema.validate(test_input)
+    res = stage_schema.validate(test_input)
 
     if expected_err == "":
         assert res.valid is True, f"err: {[e.as_dict() for e in res.errors]}"
@@ -58,14 +53,12 @@ def fake_input(tmp_path):
 
 
 @pytest.mark.skipif(not has_executable("zstd"), reason="no zstd executable")
-def test_zstd_integration(tmp_path, fake_input_tree):  # pylint: disable=unused-argument
+def test_zstd_integration(tmp_path, stage_module, fake_input_tree):  # pylint: disable=unused-argument
     inputs = fake_input_tree[1]
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.zstd")
-    stage = import_module_from_path("zstd_stage", stage_path)
     options = {
         "filename": "image.txt.zst",
     }
-    stage.main(inputs, tmp_path, options)
+    stage_module.main(inputs, tmp_path, options)
 
     img_path = os.path.join(tmp_path, "image.txt.zst")
     assert os.path.exists(img_path)
@@ -75,16 +68,14 @@ def test_zstd_integration(tmp_path, fake_input_tree):  # pylint: disable=unused-
 
 
 @mock.patch("subprocess.run")
-def test_zstd_cmdline(mock_run, tmp_path, fake_input_tree):
+def test_zstd_cmdline(mock_run, tmp_path, stage_module, fake_input_tree):
     fake_input_path = fake_input_tree[0]
     inputs = fake_input_tree[1]
-    stage_path = os.path.join(os.path.dirname(__file__), "../org.osbuild.zstd")
-    stage = import_module_from_path("zstd_stage", stage_path)
     filename = "image.txt.zst"
     options = {
         "filename": filename,
     }
-    stage.main(inputs, tmp_path, options)
+    stage_module.main(inputs, tmp_path, options)
 
     expected = [
         "zstd",