diff --git a/disagreement/client.py b/disagreement/client.py index f4557b5..d81c533 100644 --- a/disagreement/client.py +++ b/disagreement/client.py @@ -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| diff --git a/tests/test_client_uptime.py b/tests/test_client_uptime.py new file mode 100644 index 0000000..d66a87a --- /dev/null +++ b/tests/test_client_uptime.py @@ -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)