diff --git a/app/modules/latency.py b/app/modules/latency.py index 1b5a73d..b5f8728 100644 --- a/app/modules/latency.py +++ b/app/modules/latency.py @@ -1,3 +1,5 @@ +import urllib.parse + import niobot import time import httpx @@ -5,21 +7,26 @@ import httpx class LatencyModule(niobot.Module): @niobot.command("latency") - async def latency(self, ctx: niobot.Context): + async def latency(self, ctx: niobot.Context, homeserver: str = None): """See the bot's latency.""" latency = ctx.latency - homeserver = await niobot.resolve_homeserver(ctx.message.sender.split(":")[1]) + homeserver = homeserver or await niobot.resolve_homeserver(ctx.message.sender.split(":")[1]) + if not homeserver.startswith("http"): + homeserver = f"https://{homeserver}" + + parsed = urllib.parse.urlparse(homeserver) + loc = "https://" + parsed.netloc async with httpx.AsyncClient(headers={"User-Agent": niobot.__user_agent__}) as client: timings = [] for rnd in range(5): start = time.time() - await client.get(f"{homeserver}/_matrix/client/v1/ping") + await client.get(f"{loc}/_matrix/client/v1/ping") end = time.time() timings.append(end - start) fed_latency = sum(timings) / len(timings) * 1000 - hs_target = ctx.message.sender.split(":")[1] + hs_target = parsed.hostname return await ctx.respond( "Latency: {:,.2f}ms (federation latency to {!r}: {:,.2f}ms)".format(latency, hs_target, fed_latency) )