From 2769dedecafbda8497cec14b59aa510249d7f613 Mon Sep 17 00:00:00 2001
From: Niels Bantilan <niels.bantilan@gmail.com>
Date: Tue, 11 Apr 2023 13:15:57 -0400
Subject: [PATCH] fix GE and pandera tests

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>
---
 flytekit/core/base_task.py                    |  2 +-
 flytekit/core/workflow.py                     |  5 +++--
 plugins/flytekit-pandera/tests/test_plugin.py | 17 +++++++++++++++--
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py
index c2ec5f0190..829a4a1332 100644
--- a/flytekit/core/base_task.py
+++ b/flytekit/core/base_task.py
@@ -527,7 +527,7 @@ def dispatch_execute(
             except Exception as exc:
                 msg = f"Failed to convert inputs of task '{self.name}':\n  {exc}"
                 logger.error(msg)
-                raise TypeError(msg) from exc
+                raise type(exc)(msg) from exc
 
             # TODO: Logger should auto inject the current context information to indicate if the task is running within
             #   a workflow or a subworkflow etc
diff --git a/flytekit/core/workflow.py b/flytekit/core/workflow.py
index da34b98229..465d74cd66 100644
--- a/flytekit/core/workflow.py
+++ b/flytekit/core/workflow.py
@@ -262,8 +262,9 @@ def __call__(self, *args, **kwargs) -> Union[Tuple[Promise], Promise, VoidPromis
         self.compile()
         try:
             return flyte_entity_call_handler(self, *args, **input_kwargs)
-        except TypeError as exc:
-            raise TypeError(f"Encountered error while executing workflow '{self.name}':\n  {exc}") from exc
+        except Exception as exc:
+            exc.args = (f"Encountered error while executing workflow '{self.name}':\n  {exc}", *exc.args[1:])
+            raise exc
 
     def execute(self, **kwargs):
         raise Exception("Should not be called")
diff --git a/plugins/flytekit-pandera/tests/test_plugin.py b/plugins/flytekit-pandera/tests/test_plugin.py
index cc9b26c4fa..3aa9bb4d6d 100644
--- a/plugins/flytekit-pandera/tests/test_plugin.py
+++ b/plugins/flytekit-pandera/tests/test_plugin.py
@@ -55,7 +55,13 @@ def invalid_wf() -> pandera.typing.DataFrame[OutSchema]:
     def wf_with_df_input(df: pandera.typing.DataFrame[InSchema]) -> pandera.typing.DataFrame[OutSchema]:
         return transform2(df=transform1(df=df))
 
-    with pytest.raises(pandera.errors.SchemaError, match="^expected series 'col2' to have type float64, got object"):
+    with pytest.raises(
+        pandera.errors.SchemaError,
+        match=(
+            "^Encountered error while executing workflow 'test_plugin.wf_with_df_input':\n"
+            "  expected series 'col2' to have type float64, got object"
+        )
+    ):
         wf_with_df_input(df=invalid_df)
 
     # raise error when executing workflow with invalid output
@@ -67,7 +73,14 @@ def transform2_noop(df: pandera.typing.DataFrame[IntermediateSchema]) -> pandera
     def wf_invalid_output(df: pandera.typing.DataFrame[InSchema]) -> pandera.typing.DataFrame[OutSchema]:
         return transform2_noop(df=transform1(df=df))
 
-    with pytest.raises(TypeError, match="^Failed to convert return value"):
+    with pytest.raises(
+        TypeError,
+        match=(
+            "^Encountered error while executing workflow 'test_plugin.wf_invalid_output':\n"
+            "  Error encountered while executing 'wf_invalid_output':\n"
+            "  Failed to convert outputs of task"
+        ),
+    ):
         wf_invalid_output(df=valid_df)