feat(client): track connection time (#84)
This commit is contained in:
parent
c811e2b578
commit
f24c1befac
@ -18,6 +18,8 @@ from typing import (
|
||||
)
|
||||
from types import ModuleType
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from .http import HTTPClient
|
||||
from .gateway import GatewayClient
|
||||
from .shard_manager import ShardManager
|
||||
@ -36,6 +38,7 @@ from .interactions import Interaction, Snowflake
|
||||
from .error_handler import setup_global_error_handler
|
||||
from .voice_client import VoiceClient
|
||||
from .models import Activity
|
||||
from .utils import utcnow
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .models import (
|
||||
@ -168,6 +171,8 @@ class Client:
|
||||
None # The bot's own user object, populated on READY
|
||||
)
|
||||
|
||||
self.start_time: Optional[datetime] = None
|
||||
|
||||
# Internal Caches
|
||||
self._guilds: GuildCache = GuildCache()
|
||||
self._channels: ChannelCache = ChannelCache()
|
||||
@ -244,6 +249,7 @@ class Client:
|
||||
f"Client connected using {self.shard_count} shards, waiting for READY signal..."
|
||||
)
|
||||
await self.wait_until_ready()
|
||||
self.start_time = utcnow()
|
||||
print("Client is READY!")
|
||||
return
|
||||
|
||||
@ -260,6 +266,7 @@ class Client:
|
||||
# and its READY handler will set self._ready_event via dispatcher.
|
||||
print("Client connected to Gateway, waiting for READY signal...")
|
||||
await self.wait_until_ready() # Wait for the READY event from Gateway
|
||||
self.start_time = utcnow()
|
||||
print("Client is READY!")
|
||||
return # Successfully connected and ready
|
||||
except AuthenticationError: # Non-recoverable by retry here
|
||||
@ -373,6 +380,7 @@ class Client:
|
||||
await self._http.close()
|
||||
|
||||
self._ready_event.set() # Ensure any waiters for ready are unblocked
|
||||
self.start_time = None
|
||||
print("Client closed.")
|
||||
|
||||
async def __aenter__(self) -> "Client":
|
||||
@ -421,6 +429,12 @@ class Client:
|
||||
latency = getattr(self._gateway, "latency_ms", None)
|
||||
return round(latency, 2) if latency is not None else None
|
||||
|
||||
def uptime(self) -> Optional[timedelta]:
|
||||
"""Return the duration since the client connected, or ``None`` if not connected."""
|
||||
if self.start_time is None:
|
||||
return None
|
||||
return utcnow() - self.start_time
|
||||
|
||||
async def wait_until_ready(self) -> None:
|
||||
"""|coro|
|
||||
Waits until the client is fully connected to Discord and the initial state is processed.
|
||||
|
42
tests/test_client_uptime.py
Normal file
42
tests/test_client_uptime.py
Normal file
@ -0,0 +1,42 @@
|
||||
import pytest
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from disagreement.client import Client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_records_start_time(monkeypatch):
|
||||
start = datetime(2020, 1, 1, tzinfo=timezone.utc)
|
||||
|
||||
monkeypatch.setattr("disagreement.client.utcnow", lambda: start)
|
||||
|
||||
client = Client(token="t")
|
||||
monkeypatch.setattr(client, "_initialize_gateway", AsyncMock())
|
||||
client._gateway = SimpleNamespace(connect=AsyncMock())
|
||||
monkeypatch.setattr(client, "wait_until_ready", AsyncMock())
|
||||
|
||||
assert client.start_time is None
|
||||
await client.connect()
|
||||
assert client.start_time == start
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_uptime(monkeypatch):
|
||||
start = datetime(2020, 1, 1, tzinfo=timezone.utc)
|
||||
end = start + timedelta(seconds=5)
|
||||
times = [start, end]
|
||||
|
||||
def fake_now():
|
||||
return times.pop(0)
|
||||
|
||||
monkeypatch.setattr("disagreement.client.utcnow", fake_now)
|
||||
|
||||
client = Client(token="t")
|
||||
monkeypatch.setattr(client, "_initialize_gateway", AsyncMock())
|
||||
client._gateway = SimpleNamespace(connect=AsyncMock())
|
||||
monkeypatch.setattr(client, "wait_until_ready", AsyncMock())
|
||||
|
||||
await client.connect()
|
||||
assert client.uptime() == timedelta(seconds=5)
|
Loading…
x
Reference in New Issue
Block a user