diff --git a/litecord/gateway/websocket.py b/litecord/gateway/websocket.py index 6a9bf4d..613bba8 100644 --- a/litecord/gateway/websocket.py +++ b/litecord/gateway/websocket.py @@ -91,7 +91,8 @@ def _complete_users_list(user_id: str, base_ready, user_ready, ws_properties) -> for recipient in private_channel["recipients"]: users_to_send[recipient["id"]] = recipient - for relationship in user_ready["relationships"]: + user_relationships = user_ready.get("relationships", []) + for relationship in user_relationships: relationship_user = relationship["user"] users_to_send[relationship_user["id"]] = relationship_user @@ -101,7 +102,7 @@ def _complete_users_list(user_id: str, base_ready, user_ready, ws_properties) -> # relationship object structure changed in v9 if ws_properties.version == 9: ready["relationships"] = [] - for relationship in user_ready["relationships"]: + for relationship in user_relationships: ready["relationships"].append( { "user_id": relationship["user"]["id"], @@ -138,7 +139,9 @@ async def _compute_supplemental(app, base_ready, user_ready, users_to_send: dict "guilds": [], } - for relationship in user_ready["relationships"]: + user_relationships = user_ready.get("relationships", []) + + for relationship in user_relationships: if relationship["type"] != RelationshipType.FRIEND.value: continue diff --git a/tests/conftest.py b/tests/conftest.py index 417b98b..e73a1f0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -143,3 +143,25 @@ async def test_cli_staff(test_cli): yield client await client.cleanup() await _user_fixture_teardown(test_cli.app, test_user) + + +@pytest.fixture +async def test_cli_bot(test_cli): + """Yield a TestClient with a bot user.""" + # do not create a new test user to prevent race conditions caused + # by a test wanting both fixtures + app = test_cli.app + test_user = await _user_fixture_setup(app) + user_id = test_user["id"] + + assert await app.db.fetchval( + """ + UPDATE users SET bot = true WHERE id = $1 RETURNING bot + """, + user_id, + ) + + client = TestClient(test_cli, test_user) + yield client + await client.cleanup() + await _user_fixture_teardown(test_cli.app, test_user) diff --git a/tests/test_websocket.py b/tests/test_websocket.py index 403b142..166ad28 100644 --- a/tests/test_websocket.py +++ b/tests/test_websocket.py @@ -49,6 +49,28 @@ async def _close(conn): await conn.close(1000, "test end") +async def extract_and_verify_ready(conn): + ready = await _json(conn) + assert ready["op"] == OP.DISPATCH + assert ready["t"] == "READY" + + data = ready["d"] + + # NOTE: change if default gateway changes + assert data["v"] == 6 + + # make sure other fields exist and are with + # proper types. + assert isinstance(data["user"], dict) + assert isinstance(data["private_channels"], list) + assert isinstance(data["guilds"], list) + assert isinstance(data["session_id"], str) + assert isinstance(data["_trace"], list) + + if "shard" in data: + assert isinstance(data["shard"], list) + + async def get_gw(test_cli, version: int) -> str: """Get the Gateway URL.""" gw_resp = await test_cli.get(f"/api/v{version}/gateway") @@ -136,27 +158,8 @@ async def test_ready_fields(test_cli_user): ) try: - ready = await _json(conn) - assert isinstance(ready, dict) - assert ready["op"] == OP.DISPATCH - assert ready["t"] == "READY" + await extract_and_verify_ready(conn) - data = ready["d"] - assert isinstance(data, dict) - - # NOTE: change if default gateway changes - assert data["v"] == 6 - - # make sure other fields exist and are with - # proper types. - assert isinstance(data["user"], dict) - assert isinstance(data["private_channels"], list) - assert isinstance(data["guilds"], list) - assert isinstance(data["session_id"], str) - assert isinstance(data["_trace"], list) - - if "shard" in data: - assert isinstance(data["shard"], list) finally: await _close(conn) @@ -301,3 +304,17 @@ async def test_resume(test_cli_user): assert isinstance(msg, dict) assert isinstance(msg["op"], int) assert msg["op"] == OP.INVALID_SESSION + + +@pytest.mark.asyncio +async def test_ready_bot(test_cli_bot): + conn = await gw_start(test_cli_bot.cli) + await _json(conn) # ignore hello + await _json_send( + conn, {"op": OP.IDENTIFY, "d": {"token": test_cli_bot.user["token"]}} + ) + + try: + await extract_and_verify_ready(conn) + finally: + await _close(conn)