Add Error message when loading class that is not a Step

nodestream/pipeline/class_loader.py

@@ -1,4 +1,5 @@
 from importlib import import_module
+from typing import Optional, Type
 
 DECLARATIVE_INIT_METHOD_NAME = "from_file_data"
 
@@ -43,6 +44,9 @@ def find_class(class_path):
 class ClassLoader:
     """Loads a class from a string path and instantiates it with the given arguments."""
 
+    def __init__(self, class_constratint: Optional[Type] = None) -> None:
+        self.class_constratint = class_constratint or object
+
     def find_class_initializer(self, implementation, factory=None):
         class_definition = find_class(implementation)
         factory_method = factory or DECLARATIVE_INIT_METHOD_NAME
@@ -55,6 +59,12 @@ def load_class(self, implementation, arguments=None, factory=None):
         arguments = arguments or {}
         initializer = self.find_class_initializer(implementation, factory)
         try:
-            return initializer(**arguments)
+            result = initializer(**arguments)
         except TypeError as e:
             raise PipelineComponentInitializationError(initializer, arguments) from e
+
+        if not isinstance(result, self.class_constratint):
+            raise TypeError(
+                f"Expected class of type {self.class_constratint}, but got {type(result)}."
+            )
+        return result

nodestream/pipeline/pipeline_file_loader.py

@@ -9,6 +9,7 @@
 from .normalizers import Normalizer
 from .pipeline import Pipeline
 from .scope_config import ScopeConfig
+from .step import Step
 from .value_providers import ValueProvider
 
 
@@ -70,7 +71,7 @@ def for_testing(cls):
 
     def initialize_from_file_data(self, file_data: List[dict]):
         return Pipeline(
-            steps=self.load_steps(ClassLoader(), file_data),
+            steps=self.load_steps(ClassLoader(Step), file_data),
             step_outbox_size=self.step_outbox_size,
         )
 

tests/unit/pipeline/test_class_loader.py

@@ -85,3 +85,21 @@ def test_class_loader_invalid_path_invalid_arugment(subject):
             implementation="tests.unit.pipeline.test_class_loader:SimpleClass",
             arguments={"not_a_valid_argument": True},
         )
+
+
+def test_class_loader_invalid_type_constraint(subject):
+    with pytest.raises(TypeError):
+        subject.class_constratint = SimpleClassWithFactories
+        subject.load_class(
+            implementation="tests.unit.pipeline.test_class_loader:SimpleClass",
+            arguments={"argument": "test"},
+        )
+
+
+def test_class_loader_valid_subclass(subject):
+    subject.class_constratint = SimpleClass
+    subject.load_class(
+        implementation="tests.unit.pipeline.test_class_loader:SimpleClassWithFactories",
+        arguments={"argument": "test"},
+        factory="another_factory",
+    )