Repository: hikari-py/hikari
Pull request #1061 (2022-03-07 06:07:19)

Fix handling of when the first zlib message is complete
### Summary
<!-- Small summary of the merge request -->

### Checklist
<!-- Make sure to tick all the following boxes by putting an `x` in between (like this `[x]`) -->
- [ ] I have run `nox` and all the pipelines have passed.
- [ ] I have made unittests according to the code I have added/modified/deleted.

### Related issues
<!--
To mention an issue use `#issue-id` and to mention a merge request use `!merge-request-id`
To close/fix an issue use `Close #issue-id` or `Fix #issue-id` (depending on the merge request)
-->



Patch:
diff --git a/changes/1065.bugfix.md b/changes/1065.bugfix.md
new file mode 100644
index 00000000..c78b329f
--- /dev/null
+++ b/changes/1065.bugfix.md
@@ -0,0 +1,1 @@
+Add `InviteEvent` to `hikari.events.channel_events.__all__`, `hikari.events.__all__` and `hikari.__all__`.
diff --git a/changes/1066.bugfix.md b/changes/1066.bugfix.md
new file mode 100644
index 00000000..233b2a91
--- /dev/null
+++ b/changes/1066.bugfix.md
@@ -0,0 +1,1 @@
+Fix incorrect type for ATTACHMENT option values.
diff --git a/dev-requirements.txt b/dev-requirements.txt
index b7781f3d..1b5aa596 100644
--- a/dev-requirements.txt
+++ b/dev-requirements.txt
@@ -29,7 +29,7 @@ sphobjinv==2.2.1
 #################
 
 mypy==0.931
-pyright==1.1.226
+pyright==1.1.227
 
 #######################
 # DEPENDENCY CHECKING #
diff --git a/hikari/events/channel_events.py b/hikari/events/channel_events.py
index 3bc4a631..d17af655 100644
--- a/hikari/events/channel_events.py
+++ b/hikari/events/channel_events.py
@@ -37,6 +37,7 @@
     "PinsUpdateEvent",
     "GuildPinsUpdateEvent",
     "DMPinsUpdateEvent",
+    "InviteEvent",
     "InviteCreateEvent",
     "InviteDeleteEvent",
     "WebhookUpdateEvent",
diff --git a/hikari/events/interaction_events.py b/hikari/events/interaction_events.py
index 16ed8774..2a2a2cc7 100644
--- a/hikari/events/interaction_events.py
+++ b/hikari/events/interaction_events.py
@@ -23,9 +23,7 @@
 """Events fired for interaction related changes."""
 from __future__ import annotations
 
-__all__: typing.List[str] = [
-    "InteractionCreateEvent",
-]
+__all__: typing.List[str] = ["InteractionCreateEvent"]
 
 import typing
 
diff --git a/hikari/events/member_events.py b/hikari/events/member_events.py
index 5886724a..067b8f40 100644
--- a/hikari/events/member_events.py
+++ b/hikari/events/member_events.py
@@ -23,12 +23,7 @@
 """Events concerning manipulation of members within guilds."""
 from __future__ import annotations
 
-__all__: typing.List[str] = [
-    "MemberEvent",
-    "MemberCreateEvent",
-    "MemberUpdateEvent",
-    "MemberDeleteEvent",
-]
+__all__: typing.List[str] = ["MemberEvent", "MemberCreateEvent", "MemberUpdateEvent", "MemberDeleteEvent"]
 
 import abc
 import typing
diff --git a/hikari/events/role_events.py b/hikari/events/role_events.py
index 20e17a19..ca71782c 100644
--- a/hikari/events/role_events.py
+++ b/hikari/events/role_events.py
@@ -23,12 +23,7 @@
 """Events pertaining to manipulation of roles within guilds."""
 from __future__ import annotations
 
-__all__: typing.List[str] = [
-    "RoleEvent",
-    "RoleCreateEvent",
-    "RoleUpdateEvent",
-    "RoleDeleteEvent",
-]
+__all__: typing.List[str] = ["RoleEvent", "RoleCreateEvent", "RoleUpdateEvent", "RoleDeleteEvent"]
 
 import abc
 import typing
diff --git a/hikari/events/typing_events.py b/hikari/events/typing_events.py
index c497e121..6943939c 100644
--- a/hikari/events/typing_events.py
+++ b/hikari/events/typing_events.py
@@ -23,11 +23,7 @@
 """Events fired when users begin typing in channels."""
 from __future__ import annotations
 
-__all__: typing.List[str] = [
-    "TypingEvent",
-    "GuildTypingEvent",
-    "DMTypingEvent",
-]
+__all__: typing.List[str] = ["TypingEvent", "GuildTypingEvent", "DMTypingEvent"]
 
 import abc
 import typing
diff --git a/hikari/events/voice_events.py b/hikari/events/voice_events.py
index c5623ffd..ba8089ac 100644
--- a/hikari/events/voice_events.py
+++ b/hikari/events/voice_events.py
@@ -24,11 +24,7 @@
 
 from __future__ import annotations
 
-__all__: typing.List[str] = [
-    "VoiceEvent",
-    "VoiceStateUpdateEvent",
-    "VoiceServerUpdateEvent",
-]
+__all__: typing.List[str] = ["VoiceEvent", "VoiceStateUpdateEvent", "VoiceServerUpdateEvent"]
 
 import abc
 import typing
diff --git a/hikari/impl/entity_factory.py b/hikari/impl/entity_factory.py
index de26cd0b..5f1dfb30 100644
--- a/hikari/impl/entity_factory.py
+++ b/hikari/impl/entity_factory.py
@@ -72,6 +72,7 @@
     commands.OptionType.CHANNEL: snowflakes.Snowflake,
     commands.OptionType.ROLE: snowflakes.Snowflake,
     commands.OptionType.MENTIONABLE: snowflakes.Snowflake,
+    commands.OptionType.ATTACHMENT: snowflakes.Snowflake,
 }
 
 
diff --git a/hikari/impl/shard.py b/hikari/impl/shard.py
index 7516cac6..682113e9 100644
--- a/hikari/impl/shard.py
+++ b/hikari/impl/shard.py
@@ -98,6 +98,18 @@
 _CHUNKING_RATELIMIT: typing.Final[typing.Tuple[float, int]] = (60.0, 60)
 # Supported gateway version
 _VERSION: int = 8
+# Used to identify the end of a ZLIB payload
+_ZLIB_SUFFIX: typing.Final[bytes] = b"\x00\x00\xff\xff"
+# Close codes which don't invalidate the current session.
+_RECONNECTABLE_CLOSE_CODES: typing.FrozenSet[errors.ShardCloseCode] = frozenset(
+    (
+        errors.ShardCloseCode.UNKNOWN_ERROR,
+        errors.ShardCloseCode.DECODE_ERROR,
+        errors.ShardCloseCode.INVALID_SEQ,
+        errors.ShardCloseCode.SESSION_TIMEOUT,
+        errors.ShardCloseCode.RATE_LIMITED,
+    )
+)
 
 
 def _log_filterer(token: str) -> typing.Callable[[str], str]:
@@ -176,29 +188,12 @@ async def send_json(
             self.logger.log(ux.TRACE, "sending payload with size %s\n    %s", len(pl), filtered)
         await self.send_str(pl, compress)
 
-    async def _receive_and_check(self, timeout: typing.Optional[float], /) -> str:
-        message = await self.receive(timeout)
-
-        if message.type == aiohttp.WSMsgType.TEXT:
-            assert isinstance(message.data, str)
-            return message.data
-
-        if message.type == aiohttp.WSMsgType.BINARY:
-            return await self._receive_and_check_complete_zlib_package(message.data, timeout)
-
+    def _handle_other_message(self, message: aiohttp.WSMessage, /) -> typing.NoReturn:
         if message.type == aiohttp.WSMsgType.CLOSE:
             close_code = int(message.data)
             reason = message.extra
             self.logger.error("connection closed with code %s (%s)", close_code, reason)
-
-            can_reconnect = close_code < 4000 or close_code in (
-                errors.ShardCloseCode.UNKNOWN_ERROR,
-                errors.ShardCloseCode.DECODE_ERROR,
-                errors.ShardCloseCode.INVALID_SEQ,
-                errors.ShardCloseCode.SESSION_TIMEOUT,
-                errors.ShardCloseCode.RATE_LIMITED,
-            )
-
+            can_reconnect = close_code < 4000 or close_code in _RECONNECTABLE_CLOSE_CODES
             raise errors.GatewayServerClosedConnectionError(reason, close_code, can_reconnect)
 
         if message.type == aiohttp.WSMsgType.CLOSING or message.type == aiohttp.WSMsgType.CLOSED:
@@ -216,21 +211,39 @@ async def _receive_and_check(self, timeout: typing.Optional[float], /) -> str:
         )
         raise errors.GatewayError("Unexpected websocket exception from gateway") from ex
 
+    async def _receive_and_check(self, timeout: typing.Optional[float], /) -> str:
+        message = await self.receive(timeout)
+
+        if message.type == aiohttp.WSMsgType.TEXT:
+            assert isinstance(message.data, str)
+            return message.data
+
+        if message.type == aiohttp.WSMsgType.BINARY:
+            if message.data.endswith(_ZLIB_SUFFIX):
+                return self.zlib.decompress(message.data).decode("utf-8")
+
+            return await self._receive_and_check_complete_zlib_package(message.data, timeout)
+
+        self._handle_other_message(message)
+
     async def _receive_and_check_complete_zlib_package(
         self, initial_data: bytes, timeout: typing.Optional[float], /
     ) -> str:
         buff = bytearray(initial_data)
 
-        while True:
+        while not buff.endswith(_ZLIB_SUFFIX):
             message = await self.receive(timeout)
 
-            if message.type != aiohttp.WSMsgType.BINARY:
-                raise errors.GatewayError(f"Unexpected message type received {message.type.name}, expected BINARY")
+            if message.type == aiohttp.WSMsgType.BINARY:
+                buff.extend(message.data)
+                continue
+
+            if message.type == aiohttp.WSMsgType.TEXT:
+                raise errors.GatewayError("Unexpected message type received TEXT, expected BINARY")
 
-            buff.extend(message.data)
+            self._handle_other_message(message)
 
-            if buff.endswith(b"\x00\x00\xff\xff"):
-                return self.zlib.decompress(buff).decode("utf-8")
+        return self.zlib.decompress(buff).decode("utf-8")
 
     @classmethod
     @contextlib.asynccontextmanager
diff --git a/pyproject.toml b/pyproject.toml
index 19a3e06d..dcbab9a8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -79,6 +79,7 @@ norecursedirs = [
     "ci",
 ]
 filterwarnings = ["ignore:.*assertions not in test modules or plugins will be ignored because assert statements are not executed by the underlying Python interpreter.*:pytest.PytestConfigWarning"]
+required_plugins = ["pytest-asyncio"]
 
 [tool.towncrier]
 package = "hikari"


Test patch:
diff --git a/tests/hikari/impl/test_shard.py b/tests/hikari/impl/test_shard.py
index d41ffa4a..97b7fdce 100644
--- a/tests/hikari/impl/test_shard.py
+++ b/tests/hikari/impl/test_shard.py
@@ -153,56 +153,58 @@ async def test_send_json(self, transport_impl, trace):
             errors.ShardCloseCode.RATE_LIMITED,
         ],
     )
-    async def test__receive_and_check_when_message_type_is_CLOSE_and_should_reconnect(self, code, transport_impl):
+    async def test__handle_other_message_when_message_type_is_CLOSE_and_should_reconnect(self, code, transport_impl):
         stub_response = StubResponse(type=aiohttp.WSMsgType.CLOSE, extra="some error extra", data=code)
-        transport_impl.receive = mock.AsyncMock(return_value=stub_response)
 
         with pytest.raises(errors.GatewayServerClosedConnectionError) as exinfo:
-            await transport_impl._receive_and_check(10)
+            await transport_impl._handle_other_message(stub_response)
 
         exception = exinfo.value
         assert exception.reason == "some error extra"
         assert exception.code == int(code)
         assert exception.can_reconnect is True
-        transport_impl.receive.assert_awaited_once_with(10)
 
     @pytest.mark.asyncio()
     @pytest.mark.parametrize(
         "code",
         [*range(4010, 4020), 5000],
     )
-    async def test__receive_and_check_when_message_type_is_CLOSE_and_should_not_reconnect(self, code, transport_impl):
+    async def test__handle_other_message_when_message_type_is_CLOSE_and_should_not_reconnect(
+        self, code, transport_impl
+    ):
         stub_response = StubResponse(type=aiohttp.WSMsgType.CLOSE, extra="dont reconnect", data=code)
-        transport_impl.receive = mock.AsyncMock(return_value=stub_response)
 
         with pytest.raises(errors.GatewayServerClosedConnectionError) as exinfo:
-            await transport_impl._receive_and_check(10)
+            await transport_impl._handle_other_message(stub_response)
 
         exception = exinfo.value
         assert exception.reason == "dont reconnect"
         assert exception.code == int(code)
         assert exception.can_reconnect is False
-        transport_impl.receive.assert_awaited_once_with(10)
 
     @pytest.mark.asyncio()
-    async def test__receive_and_check_when_message_type_is_CLOSING(self, transport_impl):
+    async def test__handle_other_message_when_message_type_is_CLOSING(self, transport_impl):
         stub_response = StubResponse(type=aiohttp.WSMsgType.CLOSING)
-        transport_impl.receive = mock.AsyncMock(return_value=stub_response)
 
         with pytest.raises(errors.GatewayError, match="Socket has closed"):
-            await transport_impl._receive_and_check(10)
-
-        transport_impl.receive.assert_awaited_once_with(10)
+            await transport_impl._handle_other_message(stub_response)
 
     @pytest.mark.asyncio()
-    async def test__receive_and_check_when_message_type_is_CLOSED(self, transport_impl):
+    async def test__handle_other_message_when_message_type_is_CLOSED(self, transport_impl):
         stub_response = StubResponse(type=aiohttp.WSMsgType.CLOSED)
-        transport_impl.receive = mock.AsyncMock(return_value=stub_response)
 
         with pytest.raises(errors.GatewayError, match="Socket has closed"):
-            await transport_impl._receive_and_check(10)
+            await transport_impl._handle_other_message(stub_response)
 
-        transport_impl.receive.assert_awaited_once_with(10)
+    @pytest.mark.asyncio()
+    async def test__handle_other_message_when_message_type_is_unknown(self, transport_impl):
+        stub_response = mock.AsyncMock(return_value=StubResponse(type=aiohttp.WSMsgType.ERROR))
+        transport_impl.exception = mock.Mock(return_value=Exception("some error"))
+
+        with pytest.raises(errors.GatewayError, match="Unexpected websocket exception from gateway") as exc_info:
+            await transport_impl._handle_other_message(stub_response)
+
+        assert exc_info.value.__cause__ is transport_impl.exception.return_value
 
     @pytest.mark.asyncio()
     async def test__receive_and_check_when_message_type_is_BINARY(self, transport_impl):
@@ -210,14 +212,23 @@ async def test__receive_and_check_when_message_type_is_BINARY(self, transport_im
         transport_impl.receive = mock.AsyncMock(return_value=response)
         transport_impl._receive_and_check_complete_zlib_package = mock.AsyncMock()
 
-        assert (
-            await transport_impl._receive_and_check(10)
-            == transport_impl._receive_and_check_complete_zlib_package.return_value
-        )
+        result = await transport_impl._receive_and_check(10)
 
+        assert result is transport_impl._receive_and_check_complete_zlib_package.return_value
         transport_impl.receive.assert_awaited_once_with(10)
         transport_impl._receive_and_check_complete_zlib_package.assert_awaited_once_with(b"some initial data", 10)
 
+    @pytest.mark.asyncio()
+    async def test__receive_and_check_when_message_type_is_BINARY_and_the_full_payload(self, transport_impl):
+        response = StubResponse(type=aiohttp.WSMsgType.BINARY, data=b"some initial data\x00\x00\xff\xff")
+        transport_impl.receive = mock.AsyncMock(return_value=response)
+        transport_impl.zlib = mock.Mock(decompress=mock.Mock(return_value=b"aaaaaaaaaaaaaaaaaa"))
+
+        assert await transport_impl._receive_and_check(10) == "aaaaaaaaaaaaaaaaaa"
+
+        transport_impl.receive.assert_awaited_once_with(10)
+        transport_impl.zlib.decompress.assert_called_once_with(response.data)
+
     @pytest.mark.asyncio()
     async def test__receive_and_check_when_message_type_is_TEXT(self, transport_impl):
         transport_impl.receive = mock.AsyncMock(
@@ -230,16 +241,19 @@ async def test__receive_and_check_when_message_type_is_TEXT(self, transport_impl
 
     @pytest.mark.asyncio()
     async def test__receive_and_check_when_message_type_is_unknown(self, transport_impl):
+        mock_exception = errors.GatewayError("aye")
         transport_impl.receive = mock.AsyncMock(return_value=StubResponse(type=aiohttp.WSMsgType.ERROR))
-        transport_impl.exception = mock.Mock(return_value=Exception)
+        transport_impl._handle_other_message = mock.Mock(side_effect=mock_exception)
 
-        with pytest.raises(errors.GatewayError, match="Unexpected websocket exception from gateway"):
+        with pytest.raises(errors.GatewayError) as exc_info:
             await transport_impl._receive_and_check(10)
 
+        assert exc_info.value is mock_exception
         transport_impl.receive.assert_awaited_once_with(10)
+        transport_impl._handle_other_message.assert_called_once_with(transport_impl.receive.return_value)
 
     @pytest.mark.asyncio()
-    async def test__receive_and_check_complete_zlib_package_when_not_BINARY(self, transport_impl):
+    async def test__receive_and_check_complete_zlib_package_when_TEXT(self, transport_impl):
         response = StubResponse(type=aiohttp.WSMsgType.TEXT, data="not binary")
         transport_impl.receive = mock.AsyncMock(return_value=response)
 
@@ -248,6 +262,20 @@ async def test__receive_and_check_complete_zlib_package_when_not_BINARY(self, tr
 
         transport_impl.receive.assert_awaited_with(10)
 
+    @pytest.mark.asyncio()
+    async def test__receive_and_check_complete_zlib_package_for_unexpected_message_type(self, transport_impl):
+        mock_exception = errors.GatewayError("aye")
+        response = StubResponse(type=aiohttp.WSMsgType.CLOSE)
+        transport_impl.receive = mock.AsyncMock(return_value=response)
+        transport_impl._handle_other_message = mock.Mock(side_effect=mock_exception)
+
+        with pytest.raises(errors.GatewayError) as exc_info:
+            await transport_impl._receive_and_check_complete_zlib_package(b"some", 10)
+
+        assert exc_info.value is mock_exception
+        transport_impl.receive.assert_awaited_with(10)
+        transport_impl._handle_other_message.assert_called_once_with(response)
+
     @pytest.mark.asyncio()
     async def test__receive_and_check_complete_zlib_package(self, transport_impl):
         response1 = StubResponse(type=aiohttp.WSMsgType.BINARY, data=b"more")