From 50f99f52b9871bc1d2c618b6ac6cc72722ed8b54 Mon Sep 17 00:00:00 2001
From: Dat Nguyen <nltd101@gmail.com>
Date: Tue, 7 Jan 2025 10:05:53 +0700
Subject: [PATCH] add Defer Assertion

---
 examples/ExampleWithDeferAssert.py            | 16 ++++++
 src/hamcrest/core/__init__.py                 |  3 +-
 src/hamcrest/core/assert_that.py              | 18 +++++++
 tests/hamcrest_unit_test/defer_assert_test.py | 51 +++++++++++++++++++
 4 files changed, 87 insertions(+), 1 deletion(-)
 create mode 100644 examples/ExampleWithDeferAssert.py
 create mode 100644 tests/hamcrest_unit_test/defer_assert_test.py

diff --git a/examples/ExampleWithDeferAssert.py b/examples/ExampleWithDeferAssert.py
new file mode 100644
index 00000000..922c4a41
--- /dev/null
+++ b/examples/ExampleWithDeferAssert.py
@@ -0,0 +1,16 @@
+import sys
+
+sys.path.append("..")
+
+from hamcrest import *
+import unittest
+
+
+class ExampleWithDeferAssert(unittest.TestCase):
+    def testUsingDeferAssertThat(self):
+        with DeferAssertContextManager() as da:
+            da.assert_that("xx", is_("xx"))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/src/hamcrest/core/__init__.py b/src/hamcrest/core/__init__.py
index 80f11c51..0fafbfc9 100644
--- a/src/hamcrest/core/__init__.py
+++ b/src/hamcrest/core/__init__.py
@@ -1,6 +1,6 @@
 # ruff: noqa: F405, F403
 
-from hamcrest.core.assert_that import assert_that
+from hamcrest.core.assert_that import assert_that, DeferAssertContextManager
 from hamcrest.core.core import *
 
 __author__ = "Jon Reid"
@@ -23,4 +23,5 @@
     "not_none",
     "raises",
     "same_instance",
+    "DeferAssertContextManager",
 ]
diff --git a/src/hamcrest/core/assert_that.py b/src/hamcrest/core/assert_that.py
index c5940683..56deccb6 100644
--- a/src/hamcrest/core/assert_that.py
+++ b/src/hamcrest/core/assert_that.py
@@ -76,3 +76,21 @@ def _assert_bool(assertion: bool, reason: Optional[str] = None) -> None:
         if not reason:
             reason = "Assertion failed"
         raise AssertionError(reason)
+
+
+class DeferAssertContextManager:
+    def __init__(self):
+        self.exceptions = []
+
+    def assert_that(self, *args, **kwargs):
+        try:
+            assert_that(*args, **kwargs)
+        except AssertionError as e:
+            self.exceptions.append(e)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if len(self.exceptions) > 0:
+            raise self.exceptions.pop(0)
diff --git a/tests/hamcrest_unit_test/defer_assert_test.py b/tests/hamcrest_unit_test/defer_assert_test.py
new file mode 100644
index 00000000..3d9feb47
--- /dev/null
+++ b/tests/hamcrest_unit_test/defer_assert_test.py
@@ -0,0 +1,51 @@
+# encoding: utf-8
+import unittest
+
+from hamcrest.core.assert_that import DeferAssertContextManager
+from hamcrest.core.core.isequal import equal_to
+
+
+class DeferAssertContextManagerTest(unittest.TestCase):
+    def testAssertionSuccessfully(self):
+        with DeferAssertContextManager() as da:
+            da.assert_that(1, equal_to(1))
+
+    def testAssertionTeardownSuccessfully(self):
+        actual = "ACTUAL"
+
+        with DeferAssertContextManager() as da:
+            da.assert_that(actual, equal_to("ACTUAL"))
+            actual = ""
+        self.assertEqual(actual, "")
+
+    def testAssertionErrorShouldTeardownBeforeRaiseExeption(self):
+        self.maxDiff = None
+        expected = "EXPECTED"
+        actual = "ACTUAL"
+
+        expectedMessage = "\nExpected: 'EXPECTED'\n     but: was 'ACTUAL'\n"
+
+        with self.assertRaises(AssertionError) as e:
+            with DeferAssertContextManager() as da:
+                da.assert_that(actual, equal_to(expected))
+                actual = ""
+        self.assertEqual(actual, "")
+        self.assertEqual(expectedMessage, str(e.exception))
+
+    def testAssertionErrorShouldRaiseExceptionBeforeExitingContext(self):
+        self.maxDiff = None
+        expected = "EXPECTED"
+        actual = "ACTUAL"
+
+        expectedMessage = "\nExpected: 'EXPECTED'\n     but: was 'ACTUAL'\n"
+
+        with self.assertRaises(AssertionError) as e:
+            with DeferAssertContextManager() as da:
+                da.assert_that(actual, equal_to(expected))
+            actual = ""
+        self.assertNotEqual(actual, "")
+        self.assertEqual(expectedMessage, str(e.exception))
+
+
+if __name__ == "__main__":
+    unittest.main()