Define a __str__ method for FlyteException (#2203)

* define a __str__ method for FlyteException Signed-off-by: noahjax <[email protected]> * update exception __str__ and tests Signed-off-by: noahjax <[email protected]> * Fix lint errors Signed-off-by: Eduardo Apolinario <[email protected]> --------- Signed-off-by: noahjax <[email protected]> Signed-off-by: Eduardo Apolinario <[email protected]> Co-authored-by: Kevin Su <[email protected]> Co-authored-by: Eduardo Apolinario <[email protected]>
flyteorg · Mar 5, 2024 · 1c0b228 · 1c0b228
1 parent 1bb5067
commit 1c0b228
Show file tree

Hide file tree

Showing 4 changed files with 95 additions and 31 deletions.
diff --git a/flytekit/exceptions/base.py b/flytekit/exceptions/base.py
@@ -7,6 +7,13 @@ def error_code(cls):
 class FlyteException(Exception, metaclass=_FlyteCodedExceptionMetaclass):
     _ERROR_CODE = "UnknownFlyteException"
 
+    def __str__(self):
+        error_message = f"error={','.join(self.args) if self.args else 'None'}"
+        if self.__cause__:
+            error_message += f", cause={self.__cause__}"
+
+        return f"{self._ERROR_CODE}: {error_message}"
+
 
 class FlyteRecoverableException(FlyteException):
     _ERROR_CODE = "RecoverableFlyteException"
diff --git a/tests/flytekit/unit/exceptions/test_base.py b/tests/flytekit/unit/exceptions/test_base.py
@@ -5,6 +5,16 @@ def test_flyte_exception():
     try:
         raise base.FlyteException("bad")
     except Exception as e:
-        assert str(e) == "bad"
+        assert str(e) == "UnknownFlyteException: error=bad"
+        assert isinstance(type(e), base._FlyteCodedExceptionMetaclass)
+        assert type(e).error_code == "UnknownFlyteException"
+
+
+def test_flyte_exception_cause():
+    try:
+        base_exn = Exception("exception from somewhere else")
+        raise base.FlyteException("bad") from base_exn
+    except Exception as e:
+        assert str(e) == "UnknownFlyteException: error=bad, cause=exception from somewhere else"
         assert isinstance(type(e), base._FlyteCodedExceptionMetaclass)
         assert type(e).error_code == "UnknownFlyteException"
diff --git a/tests/flytekit/unit/exceptions/test_system.py b/tests/flytekit/unit/exceptions/test_system.py
@@ -3,45 +3,62 @@
 
 def test_flyte_system_exception():
     try:
-        raise system.FlyteSystemException("bad")
+        base_exn = Exception("system exploded somewhere else")
+        raise system.FlyteSystemException("bad") from base_exn
     except Exception as e:
-        assert str(e) == "bad"
+        assert str(e) == "SYSTEM:Unknown: error=bad, cause=system exploded somewhere else"
         assert isinstance(type(e), base._FlyteCodedExceptionMetaclass)
         assert type(e).error_code == "SYSTEM:Unknown"
         assert isinstance(e, base.FlyteException)
 
 
 def test_flyte_not_implemented_exception():
     try:
-        raise system.FlyteNotImplementedException("I'm lazy so I didn't implement this.")
+        base_exn = Exception("somewhere else didn't implement this")
+        raise system.FlyteNotImplementedException("I'm lazy so I didn't implement this.") from base_exn
     except Exception as e:
-        assert str(e) == "I'm lazy so I didn't implement this."
+        assert (
+            str(e)
+            == "SYSTEM:NotImplemented: error=I'm lazy so I didn't implement this., cause=somewhere else didn't implement this"
+        )
         assert isinstance(e, NotImplementedError)
         assert type(e).error_code == "SYSTEM:NotImplemented"
         assert isinstance(e, system.FlyteSystemException)
 
 
 def test_flyte_entrypoint_not_loadable_exception():
     try:
-        raise system.FlyteEntrypointNotLoadable("fake.module")
+        base_exn = Exception("somewhere else couldn't load this")
+        raise system.FlyteEntrypointNotLoadable("fake.module") from base_exn
     except Exception as e:
-        assert str(e) == "Entrypoint is not loadable!  Could not load the module: 'fake.module'."
+        assert (
+            str(e) == "SYSTEM:UnloadableCode: error=Entrypoint is not loadable!  "
+            "Could not load the module: 'fake.module'., cause=somewhere else couldn't load this"
+        )
         assert type(e).error_code == "SYSTEM:UnloadableCode"
         assert isinstance(e, system.FlyteSystemException)
 
     try:
-        raise system.FlyteEntrypointNotLoadable("fake.module", task_name="secret_task")
+        base_exn = Exception("somewhere else couldn't load this")
+        raise system.FlyteEntrypointNotLoadable("fake.module", task_name="secret_task") from base_exn
     except Exception as e:
-        assert str(e) == "Entrypoint is not loadable!  Could not find the task: 'secret_task' in 'fake.module'."
+        assert (
+            str(e) == "SYSTEM:UnloadableCode: error=Entrypoint is not loadable!  "
+            "Could not find the task: 'secret_task' in 'fake.module'., cause=somewhere else couldn't load this"
+        )
         assert type(e).error_code == "SYSTEM:UnloadableCode"
         assert isinstance(e, system.FlyteSystemException)
 
     try:
-        raise system.FlyteEntrypointNotLoadable("fake.module", additional_msg="Shouldn't have used a fake module!")
+        base_exn = Exception("somewhere else couldn't load this")
+        raise system.FlyteEntrypointNotLoadable(
+            "fake.module", additional_msg="Shouldn't have used a fake module!"
+        ) from base_exn
     except Exception as e:
         assert (
-            str(e) == "Entrypoint is not loadable!  Could not load the module: 'fake.module' "
-            "due to error: Shouldn't have used a fake module!"
+            str(e)
+            == "SYSTEM:UnloadableCode: error=Entrypoint is not loadable!  Could not load the module: 'fake.module' "
+            "due to error: Shouldn't have used a fake module!, cause=somewhere else couldn't load this"
         )
         assert type(e).error_code == "SYSTEM:UnloadableCode"
         assert isinstance(e, system.FlyteSystemException)
@@ -54,7 +71,8 @@ def test_flyte_entrypoint_not_loadable_exception():
         )
     except Exception as e:
         assert (
-            str(e) == "Entrypoint is not loadable!  Could not find the task: 'secret_task' in 'fake.module' "
+            str(e)
+            == "SYSTEM:UnloadableCode: error=Entrypoint is not loadable!  Could not find the task: 'secret_task' in 'fake.module' "
             "due to error: Shouldn't have used a fake module!"
         )
         assert type(e).error_code == "SYSTEM:UnloadableCode"
@@ -63,9 +81,13 @@ def test_flyte_entrypoint_not_loadable_exception():
 
 def test_flyte_system_assertion():
     try:
-        raise system.FlyteSystemAssertion("I assert that the system messed up.")
+        base_exn = Exception("somewhere else asserts this is wrong")
+        raise system.FlyteSystemAssertion("I assert that the system messed up.") from base_exn
     except Exception as e:
-        assert str(e) == "I assert that the system messed up."
+        assert (
+            str(e)
+            == "SYSTEM:AssertionError: error=I assert that the system messed up., cause=somewhere else asserts this is wrong"
+        )
         assert type(e).error_code == "SYSTEM:AssertionError"
         assert isinstance(e, system.FlyteSystemException)
         assert isinstance(e, AssertionError)
diff --git a/tests/flytekit/unit/exceptions/test_user.py b/tests/flytekit/unit/exceptions/test_user.py
@@ -3,81 +3,106 @@
 
 def test_flyte_user_exception():
     try:
-        raise user.FlyteUserException("bad")
+        base_exn = Exception("everywhere is bad")
+        raise user.FlyteUserException("bad") from base_exn
     except Exception as e:
-        assert str(e) == "bad"
+        assert str(e) == "USER:Unknown: error=bad, cause=everywhere is bad"
         assert isinstance(type(e), base._FlyteCodedExceptionMetaclass)
         assert type(e).error_code == "USER:Unknown"
         assert isinstance(e, base.FlyteException)
 
 
 def test_flyte_type_exception():
     try:
-        raise user.FlyteTypeException("int", "float", received_value=1, additional_msg="That was a bad idea!")
+        base_exn = Exception("Types need to be right!!!")
+        raise user.FlyteTypeException(
+            "int", "float", received_value=1, additional_msg="That was a bad idea!"
+        ) from base_exn
     except Exception as e:
-        assert str(e) == "Type error!  Received: int with value: 1, Expected: float. That was a bad idea!"
+        assert (
+            str(e) == "USER:TypeError: error=Type error!  Received: int with value: 1, Expected: float. "
+            "That was a bad idea!, cause=Types need to be right!!!"
+        )
         assert isinstance(e, TypeError)
         assert type(e).error_code == "USER:TypeError"
         assert isinstance(e, user.FlyteUserException)
 
     try:
+        base_exn = Exception("everyone is upset with you")
         raise user.FlyteTypeException(
             "int",
             ("list", "set"),
             received_value=1,
             additional_msg="That was a bad idea!",
-        )
+        ) from base_exn
     except Exception as e:
         assert (
-            str(e) == "Type error!  Received: int with value: 1, Expected one of: ('list', 'set'). That was a "
-            "bad idea!"
+            str(e)
+            == "USER:TypeError: error=Type error!  Received: int with value: 1, Expected one of: ('list', 'set'). That was a "
+            "bad idea!, cause=everyone is upset with you"
         )
         assert isinstance(e, TypeError)
         assert type(e).error_code == "USER:TypeError"
         assert isinstance(e, user.FlyteUserException)
 
     try:
-        raise user.FlyteTypeException("int", "float", additional_msg="That was a bad idea!")
+        base_exn = Exception("int != float")
+        raise user.FlyteTypeException("int", "float", additional_msg="That was a bad idea!") from base_exn
     except Exception as e:
-        assert str(e) == "Type error!  Received: int, Expected: float. That was a bad idea!"
+        assert (
+            str(e)
+            == "USER:TypeError: error=Type error!  Received: int, Expected: float. That was a bad idea!, cause=int != float"
+        )
         assert isinstance(e, TypeError)
         assert type(e).error_code == "USER:TypeError"
         assert isinstance(e, user.FlyteUserException)
 
     try:
         raise user.FlyteTypeException("int", ("list", "set"), additional_msg="That was a bad idea!")
     except Exception as e:
-        assert str(e) == "Type error!  Received: int, Expected one of: ('list', 'set'). That was a " "bad idea!"
+        assert (
+            str(e) == "USER:TypeError: error=Type error!  Received: int, Expected one of: ('list', 'set'). That was a "
+            "bad idea!"
+        )
         assert isinstance(e, TypeError)
         assert type(e).error_code == "USER:TypeError"
         assert isinstance(e, user.FlyteUserException)
 
 
 def test_flyte_value_exception():
     try:
-        raise user.FlyteValueException(-1, "Expected a value > 0")
+        base_exn = Exception("live up to expectations")
+        raise user.FlyteValueException(-1, "Expected a value > 0") from base_exn
     except user.FlyteValueException as e:
-        assert str(e) == "Value error!  Received: -1. Expected a value > 0"
+        assert (
+            str(e)
+            == "USER:ValueError: error=Value error!  Received: -1. Expected a value > 0, cause=live up to expectations"
+        )
         assert isinstance(e, ValueError)
         assert type(e).error_code == "USER:ValueError"
         assert isinstance(e, user.FlyteUserException)
 
 
 def test_flyte_assert():
     try:
-        raise user.FlyteAssertion("I ASSERT THAT THIS IS WRONG!")
+        base_exn = Exception("!!!!!")
+        raise user.FlyteAssertion("I ASSERT THAT THIS IS WRONG!") from base_exn
     except user.FlyteAssertion as e:
-        assert str(e) == "I ASSERT THAT THIS IS WRONG!"
+        assert str(e) == "USER:AssertionError: error=I ASSERT THAT THIS IS WRONG!, cause=!!!!!"
         assert isinstance(e, AssertionError)
         assert type(e).error_code == "USER:AssertionError"
         assert isinstance(e, user.FlyteUserException)
 
 
 def test_flyte_validation_error():
     try:
-        raise user.FlyteValidationException("I validated that your stuff was wrong.")
+        base_exn = Exception("somewhere else said this isn't valid")
+        raise user.FlyteValidationException("I validated that your stuff was wrong.") from base_exn
     except user.FlyteValidationException as e:
-        assert str(e) == "I validated that your stuff was wrong."
+        assert (
+            str(e)
+            == "USER:ValidationError: error=I validated that your stuff was wrong., cause=somewhere else said this isn't valid"
+        )
         assert isinstance(e, AssertionError)
         assert type(e).error_code == "USER:ValidationError"
         assert isinstance(e, user.FlyteUserException)