From aa01aa0c03411d12425703c74df45e85eae732b9 Mon Sep 17 00:00:00 2001
From: Thomas Rausch <thomas@thrau.at>
Date: Mon, 21 Aug 2023 11:04:05 +0200
Subject: [PATCH] fix Response.update_from to also copy close handlers

---
 localstack/http/response.py       |  5 +++--
 tests/integration/aws/test_app.py | 21 +++++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/localstack/http/response.py b/localstack/http/response.py
index bfbfe0721bad0..b20516cc83c59 100644
--- a/localstack/http/response.py
+++ b/localstack/http/response.py
@@ -14,13 +14,14 @@ class Response(WerkzeugResponse):
     def update_from(self, other: WerkzeugResponse):
         """
         Updates this response object with the data from the given response object. It reads the status code,
-        the response data, and updates its own headers (overwrites existing headers, but does not remove ones not
-        present in the given object).
+        the response data, and updates its own headers (overwrites existing headers, but does not remove ones
+        not present in the given object). Also updates ``call_on_close`` callbacks in the same way.
 
         :param other: the response object to read from
         """
         self.status_code = other.status_code
         self.response = other.response
+        self._on_close.extend(other._on_close)
         self.headers.update(other.headers)
 
     def set_json(self, doc: Any):
diff --git a/tests/integration/aws/test_app.py b/tests/integration/aws/test_app.py
index 8b629897ab754..e8068fc77fe1d 100644
--- a/tests/integration/aws/test_app.py
+++ b/tests/integration/aws/test_app.py
@@ -1,3 +1,5 @@
+import threading
+
 import pytest
 import requests
 import websocket
@@ -79,6 +81,25 @@ def test_404_unfortunately_detected_as_s3_request(self):
         assert "<Error><Code>NoSuchBucket</Code>" in response.text
 
 
+class TestWerkzeugIntegration:
+    def test_response_close_handlers_called_with_router(self, cleanups):
+        closed = threading.Event()
+
+        def _test_route(_request):
+            r = Response("ok", 200)
+            r.call_on_close(closed.set)
+            return r
+
+        rule = ROUTER.add("/_test/test_route", _test_route)
+        cleanups.append(lambda: ROUTER.remove(rule))
+
+        response = requests.get(get_edge_url() + "/_test/test_route")
+        assert response.status_code == 200, response.text
+        assert response.text == "ok"
+
+        assert closed.wait(timeout=3), "expected closed.set to be called"
+
+
 class TestWebSocketIntegration:
     """
     Test for the WebSocket/HandlerChain integration.