From 2e82f1f040363debe43fdde4d7dbd7acc2eb8d5a Mon Sep 17 00:00:00 2001
From: NiceAesth <nice@aesth.dev>
Date: Fri, 23 Feb 2024 15:20:40 +0200
Subject: [PATCH] feat: add resume payload

---
 pomice/models/payloads.py | 34 ++++++++++++++++++++++++++++++++++
 pomice/models/version.py  |  7 ++++++-
 pomice/pool.py            | 11 ++++++-----
 3 files changed, 46 insertions(+), 6 deletions(-)
 create mode 100644 pomice/models/payloads.py

diff --git a/pomice/models/payloads.py b/pomice/models/payloads.py
new file mode 100644
index 0000000..4ac14d6
--- /dev/null
+++ b/pomice/models/payloads.py
@@ -0,0 +1,34 @@
+from pydantic import Field
+from typing import Literal, Union
+
+from pydantic import TypeAdapter
+from pomice.models import BaseModel
+from pomice.models import LavalinkVersion
+from pomice.models import LavalinkVersion3Type
+from pomice.models import LavalinkVersion4Type
+
+__all__ = ("ResumePayload", "ResumePayloadV3", "ResumePayloadV4")
+
+
+class ResumePayload(BaseModel):
+    version: LavalinkVersion
+    timeout: int
+
+    def model_dump(self) -> dict:
+        return super().model_dump(by_alias=True, exclude={"version"})
+
+
+class ResumePayloadV3(BaseModel):
+    version: LavalinkVersion3Type
+    timeout: int
+    resuming_key: str = Field(serialization_alias="resumingKey")
+
+
+class ResumePayloadV4(BaseModel):
+    version: LavalinkVersion4Type
+    timeout: int
+    resuming: Literal[True] = True
+
+
+_ResumePayloadType = Union[ResumePayloadV3, ResumePayloadV4]
+ResumePayloadType = lambda **kwargs: TypeAdapter(_ResumePayloadType).validate_python(kwargs)
diff --git a/pomice/models/version.py b/pomice/models/version.py
index f0caaa6..8ca1ded 100644
--- a/pomice/models/version.py
+++ b/pomice/models/version.py
@@ -2,7 +2,12 @@ from typing import Literal
 from typing import NamedTuple
 from typing import Union
 
-__all__ = ("LavalinkVersion",)
+__all__ = (
+    "LavalinkVersion",
+    "LavalinkVersion3Type",
+    "LavalinkVersion4Type",
+    "LavalinkVersionType",
+)
 
 
 class LavalinkVersion(NamedTuple):
diff --git a/pomice/pool.py b/pomice/pool.py
index 6e1e0fa..5facb26 100644
--- a/pomice/pool.py
+++ b/pomice/pool.py
@@ -23,6 +23,8 @@ from discord.ext import commands
 from websockets import client
 from websockets import exceptions
 
+from pomice.models.payloads import ResumePayloadType, ResumePayloadV4
+
 from . import __version__
 from . import applemusic
 from . import spotify
@@ -300,14 +302,13 @@ class Node:
         if not self._resume_key:
             return
 
-        data = {"timeout": self._resume_timeout}
+        data = ResumePayloadType(
+            version=self._version, timeout=self._resume_timeout, resuming_key=self._resume_key
+        ).model_dump()
 
-        if self._version.major == 3:
-            data["resumingKey"] = self._resume_key
-        elif self._version.major == 4:
+        if isinstance(data, ResumePayloadV4):
             if self._log:
                 self._log.warning("Using a resume key with Lavalink v4 is deprecated.")
-            data["resuming"] = True
 
         await self.send(
             method="PATCH",