feat(speedtest): convert to pydantic

This commit is contained in:
Seaswimmer 2024-05-25 17:55:05 -04:00
parent e88cba6cf4
commit 45d0b32826
Signed by: cswimr
GPG key ID: 5D671B5D03D65A7F
3 changed files with 95 additions and 18 deletions

View file

@ -1,12 +1,14 @@
{ {
"author" : ["SeaswimmerTheFsh (seasw.)"], "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).", "install_msg" : "Thank you for installing Speedtest!\nYou can find the source code of this cog [here](https://coastalcommits.com/SeaswimmerTheFsh/SeaCogs).",
"name" : "SeaUtils", "name" : "Speedtest",
"short" : "A collection of useful utilities.", "short" : "A collection of useful utilities.",
"description" : "A collection of useful utilities.", "description" : "A collection of useful utilities.",
"end_user_data_statement" : "This cog does not store end user data.", "end_user_data_statement" : "This cog does not store end user data.",
"hidden": true, "hidden": true,
"disabled": false, "disabled": false,
"min_bot_version": "3.5.0", "min_bot_version": "3.5.0",
"min_python_version": [3, 8, 0] "min_python_version": [3, 10, 0],
"tags" : ["utility", "information"],
"requirements": ["pydantic"]
} }

72
speedtest/models.py Normal file
View file

@ -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"

View file

@ -8,13 +8,14 @@
import asyncio import asyncio
import json import json
import subprocess import subprocess
from typing import Any
import discord import discord
from redbot.core import commands from redbot.core import commands
from redbot.core.bot import Red from redbot.core.bot import Red
from redbot.core.utils import chat_formatting as cf from redbot.core.utils import chat_formatting as cf
from .models import Speedtest as sp
class Speedtest(commands.Cog): class Speedtest(commands.Cog):
"""A collection of random utilities.""" """A collection of random utilities."""
@ -35,7 +36,7 @@ class Speedtest(commands.Cog):
] ]
return "\n".join(text) return "\n".join(text)
async def run_speedtest(self) -> str | Any: async def run_speedtest(self) -> str | sp:
try: try:
process = await asyncio.create_subprocess_exec( process = await asyncio.create_subprocess_exec(
"speedtest", "-f", "json", "--accept-license", "speedtest", "-f", "json", "--accept-license",
@ -47,7 +48,7 @@ class Speedtest(commands.Cog):
stdout, stderr = await process.communicate() stdout, stderr = await process.communicate()
if process.returncode != 0: if process.returncode != 0:
return stderr.decode("utf-8") 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.command()
@commands.is_owner() @commands.is_owner()
@ -55,21 +56,23 @@ class Speedtest(commands.Cog):
"""Run a speedtest.""" """Run a speedtest."""
msg = await ctx.maybe_send_embed("Running speedtest...") msg = await ctx.maybe_send_embed("Running speedtest...")
async with ctx.typing(): async with ctx.typing():
r = await self.run_speedtest() speedtest = await self.run_speedtest()
if await ctx.embed_requested(): if await ctx.embed_requested():
if isinstance(r, str): if not isinstance(speedtest, sp):
await msg.edit(embed=discord.Embed(description=f"An error occurred! {r}", color=discord.Colour.red())) await msg.edit(embed=discord.Embed(description=f"An error occurred! {speedtest}", color=discord.Colour.red()))
return return
embed = discord.Embed(title="Speedtest Results", url=r["result"]["url"], color=await ctx.embed_color()) embed = discord.Embed(title="Speedtest Results", url=speedtest.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="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(r['ping']['latency'])} ms \nDownload: {round(r['download']['latency']['iqm'])} ms \nUpload: {round(r['upload']['latency']['iqm'])} ms") 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=r["result"]["url"] + ".png") embed.set_image(url=speedtest.result.image)
await msg.edit(embed=embed) await msg.edit(embed=embed)
else: else:
if isinstance(r, str): if not isinstance(speedtest, sp):
await msg.edit(content=f"An error occurred! {r}") await msg.edit(content=f"An error occurred! {speedtest}")
return return
await msg.edit(content=f"**Speedtest Results**\n" 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"**Bandwidth**: {speedtest.download.bandwidth / 1_000_000:.2f} Mbps down, {speedtest.download.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"**Ping**: Base: {round(speedtest.ping.latency)} ms\n"
f"**Results**: {r['result']['url']}") f"Download: {round(speedtest.download.latency.iqm)} ms\n"
f"Upload: {round(speedtest.upload.latency.iqm)} ms\n"
f"**Image**: {speedtest.result.image}")