From 7f9b8ab1c8cf162e712a49075b320d0480aeb33e Mon Sep 17 00:00:00 2001
From: Zach Probst <Zach_Probst@intuit.com>
Date: Thu, 19 Oct 2023 09:33:26 -0700
Subject: [PATCH 1/2] Add Error message when loading class that is not a Step

---
 nodestream/pipeline/class_loader.py         | 12 +++++++++++-
 nodestream/pipeline/pipeline_file_loader.py |  3 ++-
 tests/unit/pipeline/test_class_loader.py    | 18 ++++++++++++++++++
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/nodestream/pipeline/class_loader.py b/nodestream/pipeline/class_loader.py
index f1471f0df..c141169f0 100644
--- a/nodestream/pipeline/class_loader.py
+++ b/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
diff --git a/nodestream/pipeline/pipeline_file_loader.py b/nodestream/pipeline/pipeline_file_loader.py
index 70bf8e78e..15ca4cb84 100644
--- a/nodestream/pipeline/pipeline_file_loader.py
+++ b/nodestream/pipeline/pipeline_file_loader.py
@@ -8,6 +8,7 @@
 from .class_loader import ClassLoader
 from .normalizers import Normalizer
 from .pipeline import Pipeline
+from .step import Step
 from .scope_config import ScopeConfig
 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,
         )
 
diff --git a/tests/unit/pipeline/test_class_loader.py b/tests/unit/pipeline/test_class_loader.py
index 186aec517..90c363e69 100644
--- a/tests/unit/pipeline/test_class_loader.py
+++ b/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",
+    )

From a2a18204de882d1ea3146f754b7c4d2407b529df Mon Sep 17 00:00:00 2001
From: Zach Probst <zprobst@resilientvitality.com>
Date: Thu, 19 Oct 2023 16:08:02 -0700
Subject: [PATCH 2/2] Fix sort issues

---
 nodestream/pipeline/pipeline_file_loader.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nodestream/pipeline/pipeline_file_loader.py b/nodestream/pipeline/pipeline_file_loader.py
index 15ca4cb84..7a8c1984d 100644
--- a/nodestream/pipeline/pipeline_file_loader.py
+++ b/nodestream/pipeline/pipeline_file_loader.py
@@ -8,8 +8,8 @@
 from .class_loader import ClassLoader
 from .normalizers import Normalizer
 from .pipeline import Pipeline
-from .step import Step
 from .scope_config import ScopeConfig
+from .step import Step
 from .value_providers import ValueProvider