From 45d0b3282653acc5185e4618f3809e72cbf7f219 Mon Sep 17 00:00:00 2001 From: Seaswimmer Date: Sat, 25 May 2024 17:55:05 -0400 Subject: [PATCH] feat(speedtest): convert to pydantic --- speedtest/info.json | 8 +++-- speedtest/models.py | 72 ++++++++++++++++++++++++++++++++++++++++++ speedtest/speedtest.py | 33 ++++++++++--------- 3 files changed, 95 insertions(+), 18 deletions(-) create mode 100644 speedtest/models.py diff --git a/speedtest/info.json b/speedtest/info.json index f331eac..c36a53a 100644 --- a/speedtest/info.json +++ b/speedtest/info.json @@ -1,12 +1,14 @@ { "author" : ["SeaswimmerTheFsh (seasw.)"], - "install_msg" : "Thank you for installing SeaUtils!\nYou can find the source code of this cog [here](https://coastalcommits.com/SeaswimmerTheFsh/SeaCogs).", - "name" : "SeaUtils", + "install_msg" : "Thank you for installing Speedtest!\nYou can find the source code of this cog [here](https://coastalcommits.com/SeaswimmerTheFsh/SeaCogs).", + "name" : "Speedtest", "short" : "A collection of useful utilities.", "description" : "A collection of useful utilities.", "end_user_data_statement" : "This cog does not store end user data.", "hidden": true, "disabled": false, "min_bot_version": "3.5.0", - "min_python_version": [3, 8, 0] + "min_python_version": [3, 10, 0], + "tags" : ["utility", "information"], + "requirements": ["pydantic"] } diff --git a/speedtest/models.py b/speedtest/models.py new file mode 100644 index 0000000..0a60699 --- /dev/null +++ b/speedtest/models.py @@ -0,0 +1,72 @@ +from datetime import datetime + +from pydantic import BaseModel + + +class Speedtest(BaseModel): + type: str + timestamp: datetime + ping: "Ping" + download: "Bandwidth" + upload: "Bandwidth" + isp: str + interface: "Interface" + server: "Server" + result: "Result" + + @classmethod + def from_json(cls, data: dict) -> "Speedtest": + return cls( + type=data["type"], + timestamp=datetime.fromisoformat(data["timestamp"]), + ping=Ping(**data["ping"]), + download=Bandwidth(**data["download"]), + upload=Bandwidth(**data["upload"]), + isp=data["isp"], + interface=Interface(**data["interface"]), + server=Server(**data["server"]), + result=Result(**data["result"]) + ) + +class Bandwidth(BaseModel): + bandwidth: float + bytes: int + elapsed: int + latency: "Latency" + +class Interface(BaseModel): + internalIp: str + name: str + macAddr: str + isVpn: bool + externalIp: str + +class Ping(BaseModel): + jitter: float + latency: float + low: float + high: float + +class Latency(BaseModel): + iqm: float + low: float + high: float + jitter: float + +class Server(BaseModel): + id: str + name: str + location: str + country: str + host: str + port: int + ip: str + +class Result: + id: str + url: str + persisted: bool + + @property + def image(self): + return self.url + ".png" diff --git a/speedtest/speedtest.py b/speedtest/speedtest.py index a4429de..9e522a6 100644 --- a/speedtest/speedtest.py +++ b/speedtest/speedtest.py @@ -8,13 +8,14 @@ import asyncio import json import subprocess -from typing import Any import discord from redbot.core import commands from redbot.core.bot import Red from redbot.core.utils import chat_formatting as cf +from .models import Speedtest as sp + class Speedtest(commands.Cog): """A collection of random utilities.""" @@ -35,7 +36,7 @@ class Speedtest(commands.Cog): ] return "\n".join(text) - async def run_speedtest(self) -> str | Any: + async def run_speedtest(self) -> str | sp: try: process = await asyncio.create_subprocess_exec( "speedtest", "-f", "json", "--accept-license", @@ -47,7 +48,7 @@ class Speedtest(commands.Cog): stdout, stderr = await process.communicate() if process.returncode != 0: return stderr.decode("utf-8") - return json.loads(stdout.decode("utf-8")) + return sp.from_json(json.loads(stdout.decode("utf-8"))) @commands.command() @commands.is_owner() @@ -55,21 +56,23 @@ class Speedtest(commands.Cog): """Run a speedtest.""" msg = await ctx.maybe_send_embed("Running speedtest...") async with ctx.typing(): - r = await self.run_speedtest() + speedtest = await self.run_speedtest() if await ctx.embed_requested(): - if isinstance(r, str): - await msg.edit(embed=discord.Embed(description=f"An error occurred! {r}", color=discord.Colour.red())) + if not isinstance(speedtest, sp): + await msg.edit(embed=discord.Embed(description=f"An error occurred! {speedtest}", color=discord.Colour.red())) return - embed = discord.Embed(title="Speedtest Results", url=r["result"]["url"], color=await ctx.embed_color()) - embed.add_field(name="Bandwidth", value=f"{r['download']['bandwidth'] / 1_000_000:.2f} Mbps down, {r['upload']['bandwidth'] / 1_000_000:.2f} Mbps up") - embed.add_field(name="Ping", value=f"Base: {round(r['ping']['latency'])} ms \nDownload: {round(r['download']['latency']['iqm'])} ms \nUpload: {round(r['upload']['latency']['iqm'])} ms") - embed.set_image(url=r["result"]["url"] + ".png") + embed = discord.Embed(title="Speedtest Results", url=speedtest.result.url, color=await ctx.embed_color()) + embed.add_field(name="Bandwidth", value=f"{speedtest.download.bandwidth / 1_000_000:.2f} Mbps down, {speedtest.download.bandwidth / 1_000_000:.2f} Mbps up") + embed.add_field(name="Ping", value=f"Base: {round(speedtest.ping.latency)} ms \nDownload: {round(speedtest.download.latency.iqm)} ms \nUpload: {round(speedtest.upload.latency.iqm)} ms") + embed.set_image(url=speedtest.result.image) await msg.edit(embed=embed) else: - if isinstance(r, str): - await msg.edit(content=f"An error occurred! {r}") + if not isinstance(speedtest, sp): + await msg.edit(content=f"An error occurred! {speedtest}") return await msg.edit(content=f"**Speedtest Results**\n" - f"**Bandwidth**: {r['download']['bandwidth'] / 1_000_000:.2f} Mbps down, {r['upload']['bandwidth'] / 1_000_000:.2f} Mbps up\n" - f"**Ping**: Base: {round(r['ping']['latency'])} ms, Download: {round(r['download']['latency']['igm'])} ms, Upload: {round(r['upload']['latency']['igm'])} ms\n" - f"**Results**: {r['result']['url']}") + f"**Bandwidth**: {speedtest.download.bandwidth / 1_000_000:.2f} Mbps down, {speedtest.download.bandwidth / 1_000_000:.2f} Mbps up\n" + f"**Ping**: Base: {round(speedtest.ping.latency)} ms\n" + f"Download: {round(speedtest.download.latency.iqm)} ms\n" + f"Upload: {round(speedtest.upload.latency.iqm)} ms\n" + f"**Image**: {speedtest.result.image}")