feat(client): track connection time (#84)

This commit is contained in:
Slipstream 2025-06-15 15:20:06 -06:00 committed by GitHub
parent c811e2b578
commit f24c1befac
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 85 additions and 29 deletions

View File

@ -18,10 +18,12 @@ from typing import (
)
from types import ModuleType
from .http import HTTPClient
from .gateway import GatewayClient
from .shard_manager import ShardManager
from .event_dispatcher import EventDispatcher
from datetime import datetime, timedelta
from .http import HTTPClient
from .gateway import GatewayClient
from .shard_manager import ShardManager
from .event_dispatcher import EventDispatcher
from .enums import GatewayIntent, InteractionType, GatewayOpcode, VoiceRegion
from .errors import DisagreementException, AuthenticationError
from .typing import Typing
@ -35,7 +37,8 @@ from .ext import loader as ext_loader
from .interactions import Interaction, Snowflake
from .error_handler import setup_global_error_handler
from .voice_client import VoiceClient
from .models import Activity
from .models import Activity
from .utils import utcnow
if TYPE_CHECKING:
from .models import (
@ -164,9 +167,11 @@ class Client:
self._closed: bool = False
self._ready_event: asyncio.Event = asyncio.Event()
self.user: Optional["User"] = (
None # The bot's own user object, populated on READY
)
self.user: Optional["User"] = (
None # The bot's own user object, populated on READY
)
self.start_time: Optional[datetime] = None
# Internal Caches
self._guilds: GuildCache = GuildCache()
@ -239,13 +244,14 @@ class Client:
if self.shard_count and self.shard_count > 1:
await self._initialize_shard_manager()
assert self._shard_manager is not None
await self._shard_manager.start()
print(
f"Client connected using {self.shard_count} shards, waiting for READY signal..."
)
await self.wait_until_ready()
print("Client is READY!")
return
await self._shard_manager.start()
print(
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
await self._initialize_gateway()
assert self._gateway is not None # Should be initialized by now
@ -258,10 +264,11 @@ class Client:
await self._gateway.connect()
# After successful connection, GatewayClient's HELLO handler will trigger IDENTIFY/RESUME
# 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
print("Client is READY!")
return # Successfully connected and ready
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
print("Authentication failed. Please check your bot token.")
await self.close() # Ensure cleanup
@ -369,11 +376,12 @@ class Client:
if self._gateway:
await self._gateway.close()
if self._http: # HTTPClient has its own session to close
await self._http.close()
self._ready_event.set() # Ensure any waiters for ready are unblocked
print("Client closed.")
if self._http: # HTTPClient has its own session to close
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":
"""Enter the context manager by connecting to Discord."""
@ -415,11 +423,17 @@ class Client:
return self._gateway.latency
return None
@property
def latency_ms(self) -> Optional[float]:
"""Returns the gateway latency in milliseconds, or ``None`` if unavailable."""
latency = getattr(self._gateway, "latency_ms", None)
return round(latency, 2) if latency is not None else None
@property
def latency_ms(self) -> Optional[float]:
"""Returns the gateway latency in milliseconds, or ``None`` if unavailable."""
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|

View 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)