Skip to content

Commit

Permalink
Add very partial sharding
Browse files Browse the repository at this point in the history
  • Loading branch information
Not-Nik committed Jul 23, 2019
1 parent 50d8735 commit e1423b8
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 0 deletions.
6 changes: 6 additions & 0 deletions clamor/gateway/connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,10 @@ def __init__(self, url: str, **kwargs):
self._running = False
self._tg = None

# Sharding
self.shard_id = kwargs.get('shard_id')
self.shard_count = kwargs.get('shard_count')

# Heartbeat stuff
self._interval = 0
self._last_sequence = None
Expand Down Expand Up @@ -177,6 +181,8 @@ async def _identify(self):
'compress': True, # i guess?
'large_threshold': 250,
}
if self.shard_id and self.shard_count:
identify['shard'] = [self.shard_id, self.shard_count]
await self._send('IDENTIFY', identify)

async def on_open(self):
Expand Down
33 changes: 33 additions & 0 deletions tests/test_gateway_sharding.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
# -*- coding: utf-8 -*-

import unittest
import os

import anyio

from clamor import gateway, HTTP, Routes


class GatewayTests(unittest.TestCase):
def test_gateway_connect(self):
async def main():
http = HTTP(os.environ['TEST_BOT_TOKEN'])
url = await http.make_request(Routes.GET_GATEWAY)

gw_one = gateway.DiscordWebsocketClient(url['url'], shard_id=0, shard_count=2)
gw_two = gateway.DiscordWebsocketClient(url['url'], shard_id=1, shard_count=2)

self.assertIsInstance(gw_one, gateway.DiscordWebsocketClient)
self.assertIsInstance(gw_two, gateway.DiscordWebsocketClient)

async def stop_gatways(after):
await anyio.sleep(after)
await gw_one.close()
await gw_two.close()

async with anyio.create_task_group() as tg:
await tg.spawn(gw_one.start, os.environ['TEST_BOT_TOKEN'])
await tg.spawn(gw_two.start, os.environ['TEST_BOT_TOKEN'])
await tg.spawn(stop_gatways, 10)

anyio.run(main)

0 comments on commit e1423b8

Please sign in to comment.