1 files changed, 83 insertions, 4 deletions
diff --git a/tests/rest/client/test_transactions.py b/tests/rest/client/test_transactions.py
index d7cea30260..34e68ae82f 100644
--- a/tests/rest/client/test_transactions.py
+++ b/tests/rest/client/test_transactions.py
@@ -1,7 +1,11 @@
-from synapse.rest.client.transactions import HttpTransactionCache
-from synapse.rest.client.transactions import CLEANUP_PERIOD_MS
-from twisted.internet import defer
 from mock import Mock, call
+
+from twisted.internet import defer, reactor
+
+from synapse.rest.client.transactions import CLEANUP_PERIOD_MS, HttpTransactionCache
+from synapse.util import Clock
+from synapse.util.logcontext import LoggingContext
+
 from tests import unittest
 from tests.utils import MockClock
 
@@ -10,7 +14,10 @@ class HttpTransactionCacheTestCase(unittest.TestCase):
 
     def setUp(self):
         self.clock = MockClock()
-        self.cache = HttpTransactionCache(self.clock)
+        self.hs = Mock()
+        self.hs.get_clock = Mock(return_value=self.clock)
+        self.hs.get_auth = Mock()
+        self.cache = HttpTransactionCache(self.hs)
 
         self.mock_http_response = (200, "GOOD JOB!")
         self.mock_key = "foo"
@@ -40,6 +47,78 @@ class HttpTransactionCacheTestCase(unittest.TestCase):
         cb.assert_called_once_with("some_arg", keyword="arg", changing_args=0)
 
     @defer.inlineCallbacks
+    def test_logcontexts_with_async_result(self):
+        @defer.inlineCallbacks
+        def cb():
+            yield Clock(reactor).sleep(0)
+            defer.returnValue("yay")
+
+        @defer.inlineCallbacks
+        def test():
+            with LoggingContext("c") as c1:
+                res = yield self.cache.fetch_or_execute(self.mock_key, cb)
+                self.assertIs(LoggingContext.current_context(), c1)
+                self.assertEqual(res, "yay")
+
+        # run the test twice in parallel
+        d = defer.gatherResults([test(), test()])
+        self.assertIs(LoggingContext.current_context(), LoggingContext.sentinel)
+        yield d
+        self.assertIs(LoggingContext.current_context(), LoggingContext.sentinel)
+
+    @defer.inlineCallbacks
+    def test_does_not_cache_exceptions(self):
+        """Checks that, if the callback throws an exception, it is called again
+        for the next request.
+        """
+        called = [False]
+
+        def cb():
+            if called[0]:
+                # return a valid result the second time
+                return defer.succeed(self.mock_http_response)
+
+            called[0] = True
+            raise Exception("boo")
+
+        with LoggingContext("test") as test_context:
+            try:
+                yield self.cache.fetch_or_execute(self.mock_key, cb)
+            except Exception as e:
+                self.assertEqual(e.message, "boo")
+            self.assertIs(LoggingContext.current_context(), test_context)
+
+            res = yield self.cache.fetch_or_execute(self.mock_key, cb)
+            self.assertEqual(res, self.mock_http_response)
+            self.assertIs(LoggingContext.current_context(), test_context)
+
+    @defer.inlineCallbacks
+    def test_does_not_cache_failures(self):
+        """Checks that, if the callback returns a failure, it is called again
+        for the next request.
+        """
+        called = [False]
+
+        def cb():
+            if called[0]:
+                # return a valid result the second time
+                return defer.succeed(self.mock_http_response)
+
+            called[0] = True
+            return defer.fail(Exception("boo"))
+
+        with LoggingContext("test") as test_context:
+            try:
+                yield self.cache.fetch_or_execute(self.mock_key, cb)
+            except Exception as e:
+                self.assertEqual(e.message, "boo")
+            self.assertIs(LoggingContext.current_context(), test_context)
+
+            res = yield self.cache.fetch_or_execute(self.mock_key, cb)
+            self.assertEqual(res, self.mock_http_response)
+            self.assertIs(LoggingContext.current_context(), test_context)
+
+    @defer.inlineCallbacks
     def test_cleans_up(self):
         cb = Mock(
             return_value=defer.succeed(self.mock_http_response)