feat(client): track connection time (#84)
This commit is contained in:
parent
c811e2b578
commit
f24c1befac
@ -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|
|
||||
|
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