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 untrusted user: cswimr
GPG key ID: 5D671B5D03D65A7F
3 changed files with 95 additions and 18 deletions

View file

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

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 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}")