Improve type checking in _context, remove extra test

Signed-off-by: Elliot Gunton <elliotgunton@gmail.com>

src/hera/workflows/_context.py

@@ -78,9 +78,19 @@ def add_sub_node(self, node: Union[SubNodeMixin, TTemplate]) -> None:
 
         # When the user invokes a decorated function e.g. `@script inside a sub-context (dag/steps),
         # we also add the step/task's template to the overall workflow context, if it is not already added.
-        if hasattr(node, "template") and node.template is not None and not isinstance(node.template, str):
+        from hera.workflows._mixins import TemplateInvocatorSubNodeMixin
+
+        if (
+            isinstance(node, TemplateInvocatorSubNodeMixin)
+            and node.template is not None
+            and not isinstance(node.template, str)
+        ):
+            from hera.workflows.workflow import Workflow
+
+            assert isinstance(pieces[0], Workflow)
+
             found = False
-            for t in pieces[0].templates:  # type: ignore
+            for t in pieces[0].templates:
                 if t.name == node.template.name:
                     if t != node.template:
                         raise TemplateNameConflict(f"Found multiple templates with the same name: {t.name}")

tests/test_unit/test_context.py

@@ -31,6 +31,7 @@ def example():
 
     def test_no_conflict_on_tasks_with_different_names_using_same_template(self):
         """Task nodes can have different names for the same script template."""
+        dag_name = "dag-name"
         name_1 = "task-1"
         name_2 = "task-2"
 
@@ -40,11 +41,11 @@ def example():
 
         with WorkflowTemplate(
             name="my-workflow",
-            entrypoint=name,
-        ), DAG(name=name):
+            entrypoint=dag_name,
+        ), DAG(name=dag_name):
             example(name=name_1)
             example(name=name_2)
-            
+
     def test_no_conflict_on_dag_and_task_with_same_name(self):
         """Dag and task node can have the same name."""
         name = "name-of-dag-and-task"
@@ -65,34 +66,6 @@ def example():
         ), Steps(name=name):
             example(name=name)  # step name same as steps template
 
-    def test_conflict_on_multiple_scripts_with_same_name(self):
-        """Dags cannot have two scripts with the same name."""
-        name = "name-of-tasks"
-
-        @script(name=name)  # same template name
-        def hello():
-            print("hello")
-
-        @script(name=name)  # same template name
-        def world():
-            print("world")
-
-        with pytest.raises(TemplateNameConflict):
-            with WorkflowTemplate(
-                name="my-workflow",
-                entrypoint=name,
-            ), DAG(name="dag"):
-                hello()
-                world()
-
-        with pytest.raises(TemplateNameConflict):
-            with WorkflowTemplate(
-                name="my-workflow",
-                entrypoint=name,
-            ), Steps(name="steps"):
-                hello()
-                world()
-
     def test_conflict_on_multiple_tasks_with_same_name(self):
         """Dags cannot have two task nodes with the same name."""
         name = "name-of-tasks"