fix: override the token if provided in append or stop methods

Author: zimeg

slack_sdk/web/async_chat_stream.py

@@ -58,6 +58,7 @@ def __init__(
         self._buffer = ""
         self._state = "starting"
         self._stream_ts: Optional[str] = None
+        self._token: Optional[str] = kwargs.get("token")
         self._buffer_size = buffer_size
 
     async def append(
@@ -80,6 +81,8 @@ async def append(
         """
         if self._state == "completed":
             raise e.SlackRequestError(f"Cannot append to stream: stream state is {self._state}")
+        if kwargs.get("token"):
+            self._token = kwargs.pop("token")
         self._buffer += markdown_text
         if len(self._buffer) >= self._buffer_size:
             return await self._flush_buffer(**kwargs)
@@ -118,26 +121,25 @@ async def stop(
         """
         if self._state == "completed":
             raise e.SlackRequestError(f"Cannot stop stream: stream state is {self._state}")
+        if kwargs.get("token"):
+            self._token = kwargs.pop("token")
         if markdown_text:
             self._buffer += markdown_text
         if not self._stream_ts:
             response = await self._client.chat_startStream(
                 **self._stream_args,
-                **kwargs,
+                token=self._token,
             )
             if not response.get("ts"):
                 raise e.SlackRequestError("Failed to stop stream: stream not started")
             self._stream_ts = str(response["ts"])
             self._state = "in_progress"
-
-        print(f"_stream_args: {self._stream_args}\n")  # todo
-        print(f"_buffer: {self._buffer}\n")
-
         response = await self._client.chat_stopStream(
+            token=self._token,
             channel=self._stream_args["channel"],
             ts=self._stream_ts,
             blocks=blocks,
-            markdown_text=self._buffer + (markdown_text if markdown_text is not None else ""),
+            markdown_text=self._buffer,
             metadata=metadata,
             **kwargs,
         )
@@ -149,17 +151,19 @@ async def _flush_buffer(self, **kwargs) -> AsyncSlackResponse:
         if not self._stream_ts:
             response = await self._client.chat_startStream(
                 **self._stream_args,
-                markdown_text=self._buffer,
+                token=self._token,
                 **kwargs,
+                markdown_text=self._buffer,
             )
             self._stream_ts = response.get("ts")
             self._state = "in_progress"
         else:
             response = await self._client.chat_appendStream(
+                token=self._token,
                 channel=self._stream_args["channel"],
                 ts=self._stream_ts,
-                markdown_text=self._buffer,
                 **kwargs,
+                markdown_text=self._buffer,
             )
 
         self._buffer = ""

slack_sdk/web/chat_stream.py

@@ -58,6 +58,7 @@ def __init__(
         self._buffer = ""
         self._state = "starting"
         self._stream_ts: Optional[str] = None
+        self._token: Optional[str] = kwargs.get("token")
         self._buffer_size = buffer_size
 
     def append(
@@ -80,6 +81,8 @@ def append(
         """
         if self._state == "completed":
             raise e.SlackRequestError(f"Cannot append to stream: stream state is {self._state}")
+        if kwargs.get("token"):
+            self._token = kwargs.pop("token")
         self._buffer += markdown_text
         if len(self._buffer) >= self._buffer_size:
             return self._flush_buffer(**kwargs)
@@ -118,18 +121,21 @@ def stop(
         """
         if self._state == "completed":
             raise e.SlackRequestError(f"Cannot stop stream: stream state is {self._state}")
+        if kwargs.get("token"):
+            self._token = kwargs.pop("token")
         if markdown_text:
             self._buffer += markdown_text
         if not self._stream_ts:
             response = self._client.chat_startStream(
                 **self._stream_args,
-                **kwargs,
+                token=self._token,
             )
             if not response.get("ts"):
                 raise e.SlackRequestError("Failed to stop stream: stream not started")
             self._stream_ts = str(response["ts"])
             self._state = "in_progress"
         response = self._client.chat_stopStream(
+            token=self._token,
             channel=self._stream_args["channel"],
             ts=self._stream_ts,
             blocks=blocks,
@@ -145,17 +151,19 @@ def _flush_buffer(self, **kwargs) -> SlackResponse:
         if not self._stream_ts:
             response = self._client.chat_startStream(
                 **self._stream_args,
-                markdown_text=self._buffer,
+                token=self._token,
                 **kwargs,
+                markdown_text=self._buffer,
             )
             self._stream_ts = response.get("ts")
             self._state = "in_progress"
         else:
             response = self._client.chat_appendStream(
+                token=self._token,
                 channel=self._stream_args["channel"],
                 ts=self._stream_ts,
-                markdown_text=self._buffer,
                 **kwargs,
+                markdown_text=self._buffer,
             )
 
         self._buffer = ""

tests/slack_sdk/web/test_chat_stream.py

@@ -18,6 +18,7 @@ def _handle(self):
 
             # Standard auth and validation from parent
             if self.is_valid_token() and self.is_valid_user_agent():
+                token = self.headers["authorization"].split(" ")[1]
                 parsed_path = urlparse(self.path)
                 len_header = self.headers.get("Content-Length") or 0
                 content_len = int(len_header)
@@ -40,13 +41,13 @@ def _handle(self):
                 if self.path in ["/chat.startStream", "/chat.appendStream", "/chat.stopStream"] and request_body:
                     if not hasattr(self.server, "chat_stream_requests"):
                         self.server.chat_stream_requests = {}
-                    self.server.chat_stream_requests[self.path] = request_body
-
-                # Get token pattern for response file
-                header = self.headers["authorization"]
-                pattern = str(header).split("xoxb-", 1)[1]
+                    self.server.chat_stream_requests[self.path] = {
+                        "token": token,
+                        **request_body,
+                    }
 
                 # Load response file
+                pattern = str(token).split("xoxb-", 1)[1]
                 with open(f"tests/slack_sdk_fixture/web_response_{pattern}.json") as file:
                     body = json.load(file)
 
@@ -93,6 +94,7 @@ def test_streams_a_short_message(self):
         self.assertEqual(self.received_requests.get("/chat.startStream", 0), 1)
         self.assertEqual(self.received_requests.get("/chat.appendStream", 0), 0)
         self.assertEqual(self.received_requests.get("/chat.stopStream", 0), 1)
+
         if hasattr(self.thread.server, "chat_stream_requests"):
             start_request = self.thread.server.chat_stream_requests.get("/chat.startStream", {})
             self.assertEqual(start_request.get("channel"), "C0123456789")
@@ -117,7 +119,7 @@ def test_streams_a_long_message(self):
         streamer.append(markdown_text="e is")
         streamer.append(markdown_text=" bold!")
         streamer.append(markdown_text="*")
-        streamer.stop(markdown_text="*")
+        streamer.stop(markdown_text="*", token="xoxb-chat_stream_test_token")
 
         self.assertEqual(self.received_requests.get("/chat.startStream", 0), 1)
         self.assertEqual(self.received_requests.get("/chat.appendStream", 0), 1)
@@ -139,4 +141,5 @@ def test_streams_a_long_message(self):
             stop_request = self.thread.server.chat_stream_requests.get("/chat.stopStream", {})
             self.assertEqual(stop_request.get("channel"), "C0123456789")
             self.assertEqual(stop_request.get("markdown_text"), "**")
+            self.assertEqual(stop_request.get("token"), "xoxb-chat_stream_test_token")
             self.assertEqual(stop_request.get("ts"), "123.123")