disagreement/tests/test_sharding.py
2025-06-09 22:25:14 -06:00

53 lines
1.6 KiB
Python

import pytest
from unittest.mock import AsyncMock
from disagreement.shard_manager import ShardManager
from disagreement.client import Client
class DummyGateway:
def __init__(self, *args, **kwargs):
self.connect = AsyncMock()
self.close = AsyncMock()
class DummyClient:
def __init__(self):
self._http = object()
self._event_dispatcher = object()
self.token = "t"
self.intents = 0
self.verbose = False
def test_shard_manager_creates_shards(monkeypatch):
monkeypatch.setattr("disagreement.shard_manager.GatewayClient", DummyGateway)
client = DummyClient()
manager = ShardManager(client, shard_count=3)
assert len(manager.shards) == 0
manager._create_shards()
assert len(manager.shards) == 3
@pytest.mark.asyncio
async def test_shard_manager_start_and_close(monkeypatch):
monkeypatch.setattr("disagreement.shard_manager.GatewayClient", DummyGateway)
client = DummyClient()
manager = ShardManager(client, shard_count=2)
await manager.start()
for shard in manager.shards:
shard.gateway.connect.assert_awaited_once()
await manager.close()
for shard in manager.shards:
shard.gateway.close.assert_awaited_once()
@pytest.mark.asyncio
async def test_client_uses_shard_manager(monkeypatch):
dummy_manager = AsyncMock()
monkeypatch.setattr("disagreement.client.ShardManager", lambda c, n: dummy_manager)
c = Client(token="x", shard_count=2)
monkeypatch.setattr(c, "wait_until_ready", AsyncMock())
await c.connect()
dummy_manager.start.assert_awaited_once()