From 64fe38b58f8e44331144ebcb069affaaa5791556 Mon Sep 17 00:00:00 2001
From: Jorim Tielemans <tielemans.jorim@gmail.com>
Date: Tue, 28 Jan 2025 22:25:33 +0000
Subject: [PATCH] Add custom exception classes for error handling in pyrail
 package

---
 pyrail/exceptions.py | 19 +++++++++++++++++++
 pyrail/irail.py      | 16 +++++++++++-----
 2 files changed, 30 insertions(+), 5 deletions(-)
 create mode 100644 pyrail/exceptions.py

diff --git a/pyrail/exceptions.py b/pyrail/exceptions.py
new file mode 100644
index 0000000..545564f
--- /dev/null
+++ b/pyrail/exceptions.py
@@ -0,0 +1,19 @@
+"""Module containing custom exception classes for the pyrail package."""
+
+
+class RateLimitError(Exception):
+    """Raised when the rate limit is exceeded."""
+
+    pass
+
+
+class InvalidRequestError(Exception):
+    """Raised for invalid requests."""
+
+    pass
+
+
+class NotFoundError(Exception):
+    """Raised when an endpoint is not found."""
+
+    pass
diff --git a/pyrail/irail.py b/pyrail/irail.py
index 13fbdc9..16532e2 100644
--- a/pyrail/irail.py
+++ b/pyrail/irail.py
@@ -10,6 +10,7 @@
 
 from aiohttp import ClientError, ClientResponse, ClientSession
 
+from pyrail.exceptions import InvalidRequestError, NotFoundError, RateLimitError
 from pyrail.models import (
     CompositionApiResponse,
     ConnectionsApiResponse,
@@ -309,13 +310,18 @@ async def _handle_response(
             retry_after: int = int(response.headers.get("Retry-After", 1))
             logger.warning("Rate limited, retrying after %d seconds", retry_after)
             await asyncio.sleep(retry_after)
-            return await self._do_request(method, args)
+            raise RateLimitError(f"Rate limit exceeded for method {method}")
+            #return await self._do_request(method, args)
         elif response.status == 400:
-            logger.error("Bad request: %s", await response.text())
-            return None
+            error_message = await response.text()
+            logger.error("Bad request: %s", error_message)
+            raise InvalidRequestError(error_message)
+            #return None
         elif response.status == 404:
-            logger.error("Endpoint %s not found, response: %s", method, await response.text())
-            return None
+            error_message = await response.text()
+            logger.error("Endpoint %s not found, response: %s", method, error_message)
+            raise NotFoundError(error_message)
+            #return None
         elif response.status == 200:
             return await self._handle_success_response(response, method)
         elif response.status == 304: