make assert_trace work with @safe (#862) (#863)

* make assert_trace work with `@safe` (#862)

* suppress flake8 violation

* Some style fixes

* make _DesiredFunctionFound a subclass of BaseException

* fix flake8 failures

Co-authored-by: sobolevn <mail@sobolevn.me>

Author: ftruzzi

CHANGELOG.md

@@ -13,6 +13,7 @@ See [0Ver](https://0ver.org/).
 - Makes `_Nothing` a singleton
 - Fixes typos in documentation
 - Refactor `flow` function
+- Fixes that `assert_trace` was not catching containers from `@safe`-wrapped functions
 
 ## 0.15.0 aka The initial HKT release
 

returns/contrib/pytest/plugin.py

@@ -208,5 +208,5 @@ def factory(self, *args, **kwargs):
         return wraps(original)(factory)
 
 
-class _DesiredFunctionFound(RuntimeError):
+class _DesiredFunctionFound(BaseException):  # noqa: WPS418
     """Exception to raise when expected function is found."""

tests/test_contrib/test_pytest/test_plugin_has_trace.py

@@ -2,7 +2,7 @@
 
 from returns.contrib.pytest import ReturnsAsserts
 from returns.io import IOFailure, IOSuccess
-from returns.result import Failure, Success
+from returns.result import Failure, Success, safe
 
 
 def _create_container_function(container_type, container_value):
@@ -15,6 +15,12 @@ def _create_container_function_intermediate(container_type, container_value):
     )
 
 
+@safe
+def _safe_decorated_function(return_failure: bool = False):
+    if return_failure:
+        raise ValueError('Function failed')
+
+
 @pytest.mark.parametrize('container_type', [  # noqa: WPS118
     Success,
     Failure,
@@ -71,3 +77,13 @@ def test_failed_assert_trace2(
             _create_container_function_intermediate(  # type: ignore
                 wrong_type, 1,
             )
+
+
+@pytest.mark.parametrize('container_type', [  # noqa: WPS118
+    Success,
+    Failure,
+])
+def test_safe_decorated_assert(container_type, returns: ReturnsAsserts):
+    """Test if our plugin will catch containers from @safe-wrapped functions."""
+    with returns.assert_trace(container_type, _safe_decorated_function):
+        _safe_decorated_function(return_failure=container_type is Failure)