Repository: osbuild/osbuild
Pull request #561 (2021-01-15 18:03:11)

Introduce inputs and Assemblers as Stages
Inputs are new concept that provides data in various forms , like `OSTree` commits, files, or tree, which in turn can originate from either `Sources` or pipeline builds to a `Stage`. This is the last puzzle piece needed to convert `Assembler`s into `Stage`s internally and paves the way for assembler pipelines. All existing assemblers are converted to use inputs.

See individual commits for more details.

gicmo: Marking this draft because I am not 100% happy with `StoreServer.resolve_ref`. The name is iffy and maybe for now, it should also call `tree.read()`.
gicmo: > Marking this draft because I am not 100% happy with `StoreServer.resolve_ref`. The name is iffy and maybe for now, it should also call `tree.read()`.

Those both have been addressed by renaming the method to `read_tree()` and indeed calling `obj.read()` in the `StoreServer` can then keeping the context for that open for as long as `StoreServer`'s context is entered.

I am still keeping this PR as draft because another issue is worth considering: Should the input type be visible to `osbuild`, i.e. that there are two input types, like a `Source` and a `Pipeline` input (with the corresponding id). `osbuild` we need to "see" the latter in order to validate that the manifest is correct at loading time and also to determine what pipelines need to be built, because some results might be checkpointed and thus some pipelines don't need to be built. But in order to fully see the "graph" of pipelines, osbuild needs to see the pipeline ids for stage inputs.

Patch:
diff --git a/assemblers/org.osbuild.noop b/assemblers/org.osbuild.noop
index 074e8d4d..94b85dae 100755
--- a/assemblers/org.osbuild.noop
+++ b/assemblers/org.osbuild.noop
@@ -24,5 +24,7 @@ def main(_tree, _output_dir, options):
 
 if __name__ == '__main__':
     args = osbuild.api.arguments()
-    r = main(args["tree"], args["output_dir"], args.get("options", {}))
+    args_input = args["inputs"]["tree"]["path"]
+    args_output = args["tree"]
+    r = main(args_input, args_output, args.get("options", {}))
     sys.exit(r)
diff --git a/assemblers/org.osbuild.oci-archive b/assemblers/org.osbuild.oci-archive
index 30fd5a20..c514c163 100755
--- a/assemblers/org.osbuild.oci-archive
+++ b/assemblers/org.osbuild.oci-archive
@@ -272,5 +272,7 @@ def main(tree, output_dir, options):
 
 if __name__ == '__main__':
     args = osbuild.api.arguments()
-    r = main(args["tree"], args["output_dir"], args["options"])
+    args_input = args["inputs"]["tree"]["path"]
+    args_output = args["tree"]
+    r = main(args_input, args_output, args["options"])
     sys.exit(r)
diff --git a/assemblers/org.osbuild.ostree.commit b/assemblers/org.osbuild.ostree.commit
index 182cf630..b97257ef 100755
--- a/assemblers/org.osbuild.ostree.commit
+++ b/assemblers/org.osbuild.ostree.commit
@@ -185,5 +185,7 @@ def main(tree, output_dir, options, meta):
 
 if __name__ == '__main__':
     args = api.arguments()
-    r = main(args["tree"], args["output_dir"], args["options"], args["meta"])
+    args_input = args["inputs"]["tree"]["path"]
+    args_output = args["tree"]
+    r = main(args_input, args_output, args["options"], args["meta"])
     sys.exit(r)
diff --git a/assemblers/org.osbuild.qemu b/assemblers/org.osbuild.qemu
index 9d6baf3f..e6bb648b 100755
--- a/assemblers/org.osbuild.qemu
+++ b/assemblers/org.osbuild.qemu
@@ -702,5 +702,7 @@ def main(tree, output_dir, options, loop_client):
 
 if __name__ == '__main__':
     args = osbuild.api.arguments()
-    ret = main(args["tree"], args["output_dir"], args["options"], remoteloop.LoopClient("/run/osbuild/api/remoteloop"))
+    args_input = args["inputs"]["tree"]["path"]
+    args_output = args["tree"]
+    ret = main(args_input, args_output, args["options"], remoteloop.LoopClient("/run/osbuild/api/remoteloop"))
     sys.exit(ret)
diff --git a/assemblers/org.osbuild.rawfs b/assemblers/org.osbuild.rawfs
index a5d5f54f..fd4f1c97 100755
--- a/assemblers/org.osbuild.rawfs
+++ b/assemblers/org.osbuild.rawfs
@@ -107,5 +107,7 @@ def main(tree, output_dir, options, loop_client):
 
 if __name__ == '__main__':
     args = osbuild.api.arguments()
-    r = main(args["tree"], args["output_dir"], args["options"], remoteloop.LoopClient("/run/osbuild/api/remoteloop"))
+    args_input = args["inputs"]["tree"]["path"]
+    args_output = args["tree"]
+    r = main(args_input, args_output, args["options"], remoteloop.LoopClient("/run/osbuild/api/remoteloop"))
     sys.exit(r)
diff --git a/assemblers/org.osbuild.tar b/assemblers/org.osbuild.tar
index 2d8e1006..108d4a34 100755
--- a/assemblers/org.osbuild.tar
+++ b/assemblers/org.osbuild.tar
@@ -80,5 +80,7 @@ def main(tree, output_dir, options):
 
 if __name__ == '__main__':
     args = osbuild.api.arguments()
-    r = main(args["tree"], args["output_dir"], args["options"])
+    args_input = args["inputs"]["tree"]["path"]
+    args_output = args["tree"]
+    r = main(args_input, args_output, args["options"])
     sys.exit(r)
diff --git a/inputs/org.osbuild.tree b/inputs/org.osbuild.tree
new file mode 100755
index 00000000..69999c50
--- /dev/null
+++ b/inputs/org.osbuild.tree
@@ -0,0 +1,58 @@
+#!/usr/bin/python3
+"""
+Tree inputs
+
+Resolve the given pipeline `id` to a path and return that. If
+`id` is `null` or the empty string it returns an empty tree.
+"""
+
+
+import json
+import sys
+
+from osbuild.objectstore import StoreClient
+
+
+SCHEMA = """
+"additionalProperties": false,
+"required": ["pipeline"],
+"properties": {
+  "pipeline": {
+    "description": "The Pipeline that built the desired tree",
+    "type": "object",
+    "required": ["id"],
+    "additionalProperties": false,
+    "properties": {
+      "id": {
+        "description": "Identifier for the pipeline",
+        "type": "string"
+      }
+    }
+  }
+}
+"""
+
+
+def main():
+    args = json.load(sys.stdin)
+    options = args["options"]
+
+    store = StoreClient(connect_to=args["api"]["store"])
+    pid = options["pipeline"]["id"]
+
+    if not pid:
+        path = store.mkdtemp(prefix="empty")
+    else:
+        path = store.read_tree(pid)
+
+    if not path:
+        json.dump({"error": "Could find target"}, sys.stdout)
+        return 1
+
+    json.dump({"path": path}, sys.stdout)
+    return 0
+
+
+if __name__ == '__main__':
+    r = main()
+    sys.exit(r)
diff --git a/osbuild/formats/v1.py b/osbuild/formats/v1.py
index ff35d987..2413a5e0 100644
--- a/osbuild/formats/v1.py
+++ b/osbuild/formats/v1.py
@@ -2,6 +2,7 @@
 
 from typing import Dict, List, Optional, Tuple
 from osbuild.meta import Index, ValidationResult
+from ..inputs import Input
 from ..pipeline import Manifest, Pipeline, detect_host_runner
 
 
@@ -43,39 +44,45 @@ def describe(manifest: Manifest, *, with_id=False) -> Dict:
     return description
 
 
-def load_build(description: Dict, sources_options: Dict, result: List[Pipeline]):
+def load_build(description: Dict, index: Index, result: List[Pipeline]):
     pipeline = description.get("pipeline")
     if pipeline:
-        build_pipeline = load_pipeline(pipeline, sources_options, result)
+        build_pipeline = load_pipeline(pipeline, index, result)
     else:
         build_pipeline = None
 
     return build_pipeline, description["runner"]
 
 
-def load_pipeline(description: Dict, sources_options: Dict, result: List[Pipeline]) -> Pipeline:
+def load_pipeline(description: Dict, index: Index, result: List[Pipeline]) -> Pipeline:
     build = description.get("build")
     if build:
-        build_pipeline, runner = load_build(build, sources_options, result)
+        build_pipeline, runner = load_build(build, index, result)
     else:
         build_pipeline, runner = None, detect_host_runner()
 
-
-    pipeline = Pipeline(runner, build_pipeline and build_pipeline.tree_id)
+    build_id = build_pipeline and build_pipeline.tree_id
+    pipeline = Pipeline(runner, build_id)
 
     for s in description.get("stages", []):
-        pipeline.add_stage(s["name"], sources_options, s.get("options", {}))
+        info = index.get_module_info("Stage", s["name"])
+        pipeline.add_stage(info, s.get("options", {}))
 
     a = description.get("assembler")
     if a:
-        pipeline.set_assembler(a["name"], a.get("options", {}))
+        info = index.get_module_info("Assembler", a["name"])
+        asm = pipeline.set_assembler(info, a.get("options", {}))
+        info = index.get_module_info("Input", "org.osbuild.tree")
+        asm.inputs = {
+            "tree": Input(info, {"pipeline": {"id": pipeline.tree_id}})
+        }
 
     result.append(pipeline)
 
     return pipeline
 
 
-def load(description: Dict) -> Manifest:
+def load(description: Dict, index: Index) -> Manifest:
     """Load a manifest description"""
 
     pipeline = description.get("pipeline", {})
@@ -83,7 +90,11 @@ def load(description: Dict) -> Manifest:
 
     pipelines = []
 
-    load_pipeline(pipeline, sources, pipelines)
+    load_pipeline(pipeline, index, pipelines)
+
+    for pipeline in pipelines:
+        for stage in pipeline.stages:
+            stage.sources = sources
 
     manifest = Manifest(pipelines, sources)
 
diff --git a/osbuild/inputs.py b/osbuild/inputs.py
new file mode 100644
index 00000000..f0e250f1
--- /dev/null
+++ b/osbuild/inputs.py
@@ -0,0 +1,82 @@
+"""
+Pipeline inputs
+
+A pipeline input provides data in various forms to a `Stage`, like
+files, OSTree commits or trees. The content can either be obtained
+via a `Source` or have been built by a `Pipeline`. Thus an `Input`
+is the bridge between various types of content that originate from
+different types of sources.
+
+The acceptable origin of the data is determined by the `Input`
+itself. What types of input are allowed and required is determined
+by the `Stage`.
+
+To osbuild itself this is all transparent. The only data visible to
+osbuild is the path. The input options are just passed to the
+`Input` as is and the result is forwarded to the `Stage`.
+"""
+
+
+import importlib
+import json
+import os
+import subprocess
+
+from typing import Dict, Tuple
+
+from .meta import ModuleInfo
+from .objectstore import StoreServer
+
+
+class Input:
+    """
+    A single input with its corresponding options.
+    """
+
+    def __init__(self, info: ModuleInfo, options: Dict):
+        self.info = info
+        self.options = options or {}
+
+    def name(self) -> str:
+        return self.info.name
+
+    def run(self, storeapi: StoreServer) -> Tuple[str, Dict]:
+        name = self.info.name
+        msg = {
+            "options": self.options,
+            "origin": name,
+            "api": {
+                "store": storeapi.socket_address
+            }
+        }
+
+        # We want the `osbuild` python package that contains this
+        # very module, which might be different from the system wide
+        # installed one, to be accessible to the Input programs so
+        # we detect our origin and set the `PYTHONPATH` accordingly
+        modorigin = importlib.util.find_spec("osbuild").origin
+        modpath = os.path.dirname(modorigin)
+        env = os.environ.copy()
+        env["PYTHONPATH"] = os.path.dirname(modpath)
+
+        r = subprocess.run([self.info.path],
+                           env=env,
+                           input=json.dumps(msg),
+                           stdout=subprocess.PIPE,
+                           encoding="utf-8",
+                           check=False)
+
+        try:
+            reply = json.loads(r.stdout)
+        except ValueError:
+            raise RuntimeError(f"{name}: error: {r.stderr}") from None
+
+        if "error" in reply:
+            raise RuntimeError(f"{name}: " + reply["error"])
+
+        if r.returncode != 0:
+            raise RuntimeError(f"{name}: error {r.returncode}")
+
+        path, data = reply["path"], reply.get("data", {})
+
+        return path, data
diff --git a/osbuild/main_cli.py b/osbuild/main_cli.py
index 789a8ea6..c1b913de 100644
--- a/osbuild/main_cli.py
+++ b/osbuild/main_cli.py
@@ -90,7 +90,7 @@ def osbuild_cli():
             show_validation(res, args.manifest_path)
         return 2
 
-    manifest = fmt.load(desc)
+    manifest = fmt.load(desc, index)
 
     if args.checkpoint:
         missed = manifest.mark_checkpoints(args.checkpoint)
diff --git a/osbuild/meta.py b/osbuild/meta.py
index 691337b3..ea1718ef 100644
--- a/osbuild/meta.py
+++ b/osbuild/meta.py
@@ -273,9 +273,10 @@ class ModuleInfo:
     Normally this class is instantiated via its `load` method.
     """
 
-    def __init__(self, klass: str, name: str, info: str):
+    def __init__(self, klass: str, name: str, path: str, info: str):
         self.name = name
         self.type = klass
+        self.path = path
 
         opts = info.get("schema") or ""
         self.info = info.get("info")
@@ -353,14 +354,15 @@ class ModuleInfo:
             'desc': doclist[0],
             'info': "\n".join(doclist[1:])
         }
-        return cls(klass, name, info)
+        return cls(klass, name, path, info)
 
     @staticmethod
     def module_class_to_directory(klass: str) -> str:
         mapping = {
-            "Stage": "stages",
             "Assembler": "assemblers",
-            "Source": "sources"
+            "Input": "inputs",
+            "Source": "sources",
+            "Stage": "stages",
         }
 
         return mapping.get(klass)
@@ -418,7 +420,7 @@ class Index:
             with contextlib.suppress(FileNotFoundError):
                 with open(path, "r") as f:
                     schema = json.load(f)
-        elif klass in ["Stage", "Assembler", "Source"]:
+        elif klass in ["Assembler", "Input", "Source", "Stage"]:
             info = self.get_module_info(klass, name)
             if info:
                 schema = info.schema
diff --git a/osbuild/objectstore.py b/osbuild/objectstore.py
index a24e6fdc..28cbfa11 100644
--- a/osbuild/objectstore.py
+++ b/osbuild/objectstore.py
@@ -7,7 +7,8 @@ import tempfile
 from typing import Optional
 
 from osbuild.util.types import PathLike
-from osbuild.util import ctx, rmrf
+from osbuild.util import ctx, jsoncomm, rmrf
+from . import api
 from . import treesum
 
 
@@ -360,3 +361,79 @@ class ObjectStore(contextlib.AbstractContextManager):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.cleanup()
+
+
+class StoreServer(api.BaseAPI):
+
+    endpoint = "store"
+
+    def __init__(self, store: ObjectStore, *, socket_address=None):
+        super().__init__(socket_address)
+        self.store = store
+        self.tmproot = store.tempdir(prefix="store-server-")
+        self._stack = contextlib.ExitStack()
+
+    def _cleanup(self):
+        self.tmproot.cleanup()
+        self.tmproot = None
+        self._stack.close()
+        self._stack = None
+
+    def _read_tree(self, msg, sock):
+        object_id = msg["object-id"]
+        obj = self.store.get(object_id)
+        if not obj:
+            sock.send({"path": None})
+            return
+
+        reader = obj.read()
+        path = self._stack.enter_context(reader)
+        sock.send({"path": path})
+
+    def _mkdtemp(self, msg, sock):
+        args = {
+            "suffix": msg.get("suffix"),
+            "prefix": msg.get("prefix"),
+            "dir": self.tmproot.name
+        }
+
+        path = tempfile.mkdtemp(**args)
+        sock.send({"path": path})
+
+    def _message(self, msg, _fds, sock):
+        if msg["method"] == "read-tree":
+            self._read_tree(msg, sock)
+        if msg["method"] == "mkdtemp":
+            self._mkdtemp(msg, sock)
+
+
+class StoreClient:
+    def __init__(self, connect_to="/run/osbuild/api/store"):
+        self.client = jsoncomm.Socket.new_client(connect_to)
+
+    def __del__(self):
+        if self.client is not None:
+            self.client.close()
+
+    def mkdtemp(self, suffix=None, prefix=None):
+        msg = {
+            "method": "mkdtemp",
+            "suffix": suffix,
+            "prefix": prefix
+        }
+
+        self.client.send(msg)
+        msg, _, _ = self.client.recv()
+
+        return msg["path"]
+
+    def read_tree(self, object_id: str):
+        msg = {
+            "method": "read-tree",
+            "object-id": object_id
+        }
+
+        self.client.send(msg)
+        msg, _, _ = self.client.recv()
+
+        return msg["path"]
diff --git a/osbuild/pipeline.py b/osbuild/pipeline.py
index 0a5e6ee8..314d0ad8 100644
--- a/osbuild/pipeline.py
+++ b/osbuild/pipeline.py
@@ -1,8 +1,7 @@
-
+import contextlib
 import hashlib
 import json
 import os
-import tempfile
 from typing import Dict, List
 
 from .api import API
@@ -33,13 +32,18 @@ class BuildResult:
 
 
 class Stage:
-    def __init__(self, name, source_options, build, base, options):
-        self.name = name
+    def __init__(self, info, source_options, build, base, options):
+        self.info = info
         self.sources = source_options
         self.build = build
         self.base = base
         self.options = options
         self.checkpoint = False
+        self.inputs = {}
+
+    @property
+    def name(self):
+        return self.info.name
 
     @property
     def id(self):
@@ -50,38 +54,55 @@ class Stage:
         m.update(json.dumps(self.options, sort_keys=True).encode())
         return m.hexdigest()
 
-    def run(self,
-            tree,
-            runner,
-            build_tree,
-            cache,
-            monitor,
-            libdir,
-            var="/var/tmp"):
-        with buildroot.BuildRoot(build_tree, runner, libdir, var=var) as build_root, \
-            tempfile.TemporaryDirectory(prefix="osbuild-sources-output-", dir=var) as sources_output:
+    def run(self, tree, runner, build_tree, store, monitor, libdir):
+        with contextlib.ExitStack() as cm:
+
+            build_root = buildroot.BuildRoot(build_tree, runner, libdir, store.tmp)
+            cm.enter_context(build_root)
+
+            sources_tmp = store.tempdir(prefix="osbuild-sources-output-")
+            sources_output = cm.enter_context(sources_tmp)
 
             args = {
                 "tree": "/run/osbuild/tree",
                 "sources": "/run/osbuild/sources",
                 "options": self.options,
+                "inputs": {},
                 "meta": {
                     "id": self.id
                 }
             }
 
-            ro_binds = [f"{sources_output}:/run/osbuild/sources"]
+            ro_binds = [
+                f"{self.info.path}:/run/osbuild/bin/{self.name}",
+                f"{sources_output}:/run/osbuild/sources"
+            ]
+
+            storeapi = objectstore.StoreServer(store)
+            cm.enter_context(storeapi)
+
+            for key, ip in self.inputs.items():
+                path, data = ip.run(storeapi)
+
+                # bind mount the returned path into the container
+                mapped = f"/run/osbuild/inputs/{key}"
+                ro_binds += [f"{path}:{mapped}"]
+
+                args["inputs"][key] = {"path": mapped, "data": data}
 
             api = API(args, monitor)
             build_root.register_api(api)
 
             src = sources.SourcesServer(libdir,
                                         self.sources,
-                                        os.path.join(cache, "sources"),
+                                        os.path.join(store.store, "sources"),
                                         sources_output)
             build_root.register_api(src)
 
-            r = build_root.run([f"/run/osbuild/lib/stages/{self.name}"],
+            rls = remoteloop.LoopServer()
+            build_root.register_api(rls)
+
+            r = build_root.run([f"/run/osbuild/bin/{self.name}"],
                                monitor,
                                binds=[os.fspath(tree) + ":/run/osbuild/tree"],
                                readonly_binds=ro_binds)
@@ -159,14 +180,16 @@ class Pipeline:
     def output_id(self):
         return self.assembler.id if self.assembler else None
 
-    def add_stage(self, name, sources_options=None, options=None):
-        stage = Stage(name, sources_options, self.build, self.tree_id, options or {})
+    def add_stage(self, info, options, sources_options=None):
+        stage = Stage(info, sources_options, self.build, self.tree_id, options or {})
         self.stages.append(stage)
         if self.assembler:
             self.assembler.base = stage.id
+        return stage
 
-    def set_assembler(self, name, options=None):
-        self.assembler = Assembler(name, self.build, self.tree_id, options or {})
+    def set_assembler(self, info, options=None):
+        self.assembler = Stage(info, {}, self.build, self.tree_id, options or {})
+        return self.assembler
 
     def build_stages(self, object_store, monitor, libdir):
         results = {"success": True}
@@ -224,10 +247,9 @@ class Pipeline:
                 r = stage.run(path,
                               self.runner,
                               build_path,
-                              object_store.store,
+                              object_store,
                               monitor,
-                              libdir,
-                              var=object_store.store)
+                              libdir)
 
                 monitor.result(r)
 
@@ -246,7 +268,7 @@ class Pipeline:
 
         return results, build_tree, tree
 
-    def assemble(self, object_store, build_tree, tree, monitor, libdir):
+    def assemble(self, object_store, build_tree, monitor, libdir):
         results = {"success": True}
 
         if not self.assembler:
@@ -255,18 +277,16 @@ class Pipeline:
         output = object_store.new()
 
         with build_tree.read() as build_dir, \
-             tree.read() as input_dir, \
              output.write() as output_dir:
 
             monitor.assembler(self.assembler)
 
-            r = self.assembler.run(input_dir,
+            r = self.assembler.run(output_dir,
                                    self.runner,
                                    build_dir,
+                                   object_store,
                                    monitor,
-                                   libdir,
-                                   output_dir,
-                                   var=object_store.store)
+                                   libdir)
 
             monitor.result(r)
 
@@ -294,16 +314,13 @@ class Pipeline:
         obj = store.get(self.output_id)
 
         if not obj:
-            results, build_tree, tree = self.build_stages(store,
-                                                          monitor,
-                                                          libdir)
+            results, build_tree, _ = self.build_stages(store, monitor, libdir)
 
             if not results["success"]:
                 return results
 
             r, obj = self.assemble(store,
                                    build_tree,
-                                   tree,
                                    monitor,
                                    libdir)
 


Test patch:
diff --git a/test/mod/test_monitor.py b/test/mod/test_monitor.py
index 82593ad1..5f3e93c7 100644
--- a/test/mod/test_monitor.py
+++ b/test/mod/test_monitor.py
@@ -55,12 +55,16 @@ class TestMonitor(unittest.TestCase):
     @unittest.skipUnless(test.TestBase.can_bind_mount(), "root-only")
     def test_log_monitor_vfuncs(self):
         # Checks the basic functioning of the LogMonitor
+        index = osbuild.meta.Index(os.curdir)
+
         runner = detect_host_runner()
         pipeline = osbuild.Pipeline(runner=runner)
-        pipeline.add_stage("org.osbuild.noop", {}, {
+        info = index.get_module_info("Stage", "org.osbuild.noop")
+        pipeline.add_stage(info, {
             "isthisthereallife": False
         })
-        pipeline.set_assembler("org.osbuild.noop")
+        info = index.get_module_info("Assembler", "org.osbuild.noop")
+        pipeline.set_assembler(info)
 
         with tempfile.TemporaryDirectory() as tmpdir:
             storedir = os.path.join(tmpdir, "store")
@@ -87,14 +91,18 @@ class TestMonitor(unittest.TestCase):
     def test_monitor_integration(self):
         # Checks the monitoring API is called properly from the pipeline
         runner = detect_host_runner()
+        index = osbuild.meta.Index(os.curdir)
+
         pipeline = osbuild.Pipeline(runner=runner)
-        pipeline.add_stage("org.osbuild.noop", {}, {
+        noop_info = index.get_module_info("Stage", "org.osbuild.noop")
+        pipeline.add_stage(noop_info, {
             "isthisthereallife": False
         })
-        pipeline.add_stage("org.osbuild.noop", {}, {
+        pipeline.add_stage(noop_info, {
             "isthisjustfantasy": True
         })
-        pipeline.set_assembler("org.osbuild.noop")
+        info = index.get_module_info("Assembler", "org.osbuild.noop")
+        pipeline.set_assembler(info)
 
         with tempfile.TemporaryDirectory() as tmpdir:
             storedir = os.path.join(tmpdir, "store")
diff --git a/test/mod/test_osbuild.py b/test/mod/test_osbuild.py
index 9148ead6..b1db3757 100644
--- a/test/mod/test_osbuild.py
+++ b/test/mod/test_osbuild.py
@@ -13,6 +13,7 @@ import osbuild
 import osbuild.meta
 from osbuild.formats import v1 as fmt
 from osbuild.monitor import NullMonitor
+from osbuild.objectstore import ObjectStore
 from osbuild.pipeline import detect_host_runner
 from .. import test
 
@@ -22,6 +23,8 @@ class TestDescriptions(unittest.TestCase):
         """Degenerate case. Make sure we always return the same canonical
         description when passing empty or null values."""
 
+        index = osbuild.meta.Index(os.curdir)
+
         cases = [
             {},
             {"assembler": None},
@@ -32,59 +35,42 @@ class TestDescriptions(unittest.TestCase):
         for pipeline in cases:
             manifest = {"pipeline": pipeline}
             with self.subTest(pipeline):
-                desc = fmt.describe(fmt.load(manifest))
+                desc = fmt.describe(fmt.load(manifest, index))
                 self.assertEqual(desc["pipeline"], {})
 
     @unittest.skipUnless(test.TestBase.can_bind_mount(), "root-only")
     def test_stage_run(self):
-        stage = osbuild.Stage("org.osbuild.noop", {}, None, None, {})
+        index = osbuild.meta.Index(os.curdir)
+        info = index.get_module_info("Stage", "org.osbuild.noop")
+        stage = osbuild.Stage(info, {}, None, None, {})
 
         with tempfile.TemporaryDirectory() as tmpdir:
 
             data = pathlib.Path(tmpdir, "data")
-            cache = pathlib.Path(tmpdir, "cache")
+            storedir = pathlib.Path(tmpdir, "store")
             root = pathlib.Path("/")
             runner = detect_host_runner()
             monitor = NullMonitor(sys.stderr.fileno())
             libdir = os.path.abspath(os.curdir)
+            store = ObjectStore(storedir)
+            data.mkdir()
 
-            for p in [data, cache]:
-                p.mkdir()
-
-            res = stage.run(data, runner, root, cache, monitor, libdir)
+            res = stage.run(data, runner, root, store, monitor, libdir)
 
         self.assertEqual(res.success, True)
         self.assertEqual(res.id, stage.id)
 
-    @unittest.skipUnless(test.TestBase.can_bind_mount(), "root-only")
-    def test_assembler_run(self):
-        asm = osbuild.Assembler("org.osbuild.noop", None, None, {})
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-
-            data = pathlib.Path(tmpdir, "data")
-            cache = pathlib.Path(tmpdir, "cache")
-            output = pathlib.Path(tmpdir, "output")
-            root = pathlib.Path("/")
-            runner = detect_host_runner()
-            monitor = NullMonitor(sys.stderr.fileno())
-            libdir = os.path.abspath(os.curdir)
-
-            for p in [data, cache, output]:
-                p.mkdir()
-
-            res = asm.run(data, runner, root, monitor, libdir, output)
-
-        self.assertEqual(res.success, True)
-        self.assertEqual(res.id, asm.id)
-
     def test_pipeline(self):
+        index = osbuild.meta.Index(os.curdir)
+
+        test_info = index.get_module_info("Stage", "org.osbuild.test")
         build = osbuild.Pipeline("org.osbuild.test")
-        build.add_stage("org.osbuild.test", {}, {"one": 1})
+        build.add_stage(test_info, {"one": 1})
 
         pipeline = osbuild.Pipeline("org.osbuild.test", build.tree_id)
-        pipeline.add_stage("org.osbuild.test", {}, {"one": 2})
-        pipeline.set_assembler("org.osbuild.test")
+        pipeline.add_stage(test_info, {"one": 2})
+        info = index.get_module_info("Assembler", "org.osbuild.noop")
+        pipeline.set_assembler(info)
 
         manifest = osbuild.Manifest([build, pipeline], {})
 
@@ -108,7 +94,7 @@ class TestDescriptions(unittest.TestCase):
                     }
                 ],
                 "assembler": {
-                    "name": "org.osbuild.test"
+                    "name": "org.osbuild.noop"
                 }
             }
         })
@@ -117,7 +103,7 @@ class TestDescriptions(unittest.TestCase):
         index = osbuild.meta.Index(os.curdir)
 
         modules = []
-        for klass in ("Stage", "Assembler", "Source"):
+        for klass in ("Assembler", "Input", "Source", "Stage"):
             mods = index.list_modules_for_class(klass)
             modules += [(klass, module) for module in mods]
 
diff --git a/test/test.py b/test/test.py
index 9a2e1a8b..755d377b 100644
--- a/test/test.py
+++ b/test/test.py
@@ -11,6 +11,7 @@ import sys
 import tempfile
 import unittest
 
+import osbuild.meta
 from osbuild.formats import v1 as fmt
 from osbuild.util import linux
 
@@ -352,9 +353,11 @@ class OSBuild(contextlib.AbstractContextManager):
         are defined.
         """
 
+        index = osbuild.meta.Index(os.curdir)
+
         manifest_json = json.loads(manifest_data)
 
-        manifest = fmt.load(manifest_json)
+        manifest = fmt.load(manifest_json, index)
         tree_id, _ = fmt.get_ids(manifest)
         return tree_id