Added tests

Author: eadwinCode

ellar_throttler/__init__.py

@@ -1,3 +1,21 @@
 """A rate limiting module for Ellar"""
 
 __version__ = "0.0.1"
+
+from .decorators import skip_throttle, throttle
+from .exception import ThrottledException
+from .interfaces import IThrottlerStorage
+from .module import ThrottlerModule
+from .throttler_guard import ThrottlerGuard
+from .throttler_service import CacheThrottlerStorageService, ThrottlerStorageService
+
+__all__ = [
+    "throttle",
+    "skip_throttle",
+    "ThrottlerModule",
+    "ThrottlerGuard",
+    "ThrottledException",
+    "CacheThrottlerStorageService",
+    "ThrottlerStorageService",
+    "IThrottlerStorage",
+]

ellar_throttler/throttler_service.py

@@ -15,20 +15,21 @@ class ThrottlerStorageService(IThrottlerStorage):
     def __init__(self) -> None:
         self._storage: t.Dict[str, ThrottlerStorageOption] = {}
 
-    def _set(self, key: str, value: ThrottlerStorageOption) -> None:
-        self._storage[key] = value
+    def _set(self, key: str, value: ThrottlerStorageOption) -> ThrottlerStorageOption:
+        self.storage[key] = value
+        return value
 
     def _has_expired(self, key: str) -> bool:
-        exp = self._storage.get(key)
+        exp = self.storage.get(key)
         return exp is not None and exp.expires_at <= time.time()
 
     def _has_key(self, key: str) -> bool:
         return self._get(key) is not None
 
     def _delete(self, key: str) -> bool:
         try:
-            self._storage.pop(key)
-        except KeyError:
+            self.storage.pop(key)
+        except KeyError:  # pragma: no cover
             return False
         return True
 
@@ -37,14 +38,17 @@ def _get(self, key: str) -> t.Optional[ThrottlerStorageOption]:
             self._delete(key)
             return None
 
-        return self._storage.get(key)
+        return self.storage.get(key)
 
     @property
     def storage(self) -> t.Dict[str, ThrottlerStorageOption]:
         return self._storage
 
     def get_expiration_time(self, key: str) -> int:
-        return math.floor(self.storage[key].expires_at - time.time())
+        cache = self._get(key)
+        if cache is None:  # pragma: no cover
+            return -1
+        return math.floor(cache.expires_at - time.time())
 
     async def increment(self, key: str, ttl: int) -> ThrottlerStorageRecord:
         if not self._has_key(key):
@@ -53,20 +57,20 @@ async def increment(self, key: str, ttl: int) -> ThrottlerStorageRecord:
             )
 
         history = self._get(key)
-        assert history, "value can not be None"
-        now = time.time()
 
         time_to_expire = self.get_expiration_time(key)
 
         if time_to_expire <= 0:
-            history.total_hits = 0
-            history.expires_at = now + ttl
+            history = self._set(
+                key, ThrottlerStorageOption(total_hits=0, expires_at=time.time() + ttl)
+            )
             time_to_expire = self.get_expiration_time(key)
 
-        history.total_hits += 1
+        history.total_hits += 1  # type:ignore[union-attr]
 
         return ThrottlerStorageRecord(
-            total_hits=history.total_hits, time_to_expire=time_to_expire
+            total_hits=history.total_hits,  # type:ignore[union-attr]
+            time_to_expire=time_to_expire,
         )
 
 
@@ -76,7 +80,7 @@ def __init__(self, cache_service: ICacheService) -> None:
         self._cache_service = cache_service
 
     @property
-    def storage(self) -> t.Any:
+    def storage(self) -> t.Any:  # pragma: no cover
         return self._cache_service.get_backend()
 
     async def get_expiration_time(self, key: str) -> int:

tests/app/__init__.py

@@ -0,0 +1,21 @@
+from ellar.common import Controller, get
+
+from ellar_throttler import skip_throttle, throttle
+
+from .service import AppService
+
+
+@Controller("/")
+@throttle(limit=2, ttl=10)
+class AppController:
+    def __init__(self, app_service: AppService):
+        self.app_service = app_service
+
+    @get()
+    async def test(self):
+        return self.app_service.success()
+
+    @get("/ignored")
+    @skip_throttle()
+    async def ignored(self):
+        return self.app_service.ignored()

tests/app/controller/__init__.py

@@ -0,0 +1,13 @@
+from ellar.common import Controller, get
+
+from .service import AppService
+
+
+@Controller("/default")
+class DefaultController:
+    def __init__(self, app_service: AppService):
+        self.app_service = app_service
+
+    @get()
+    def get_default(self):
+        return self.app_service.success()

tests/app/controller/app.py

@@ -0,0 +1,31 @@
+from ellar.common import Controller, get
+
+from ellar_throttler import throttle
+
+from .service import AppService
+
+
+@throttle(limit=2, ttl=10)
+@Controller("/limit")
+class LimitController:
+    def __init__(self, app_service: AppService):
+        self.app_service = app_service
+
+    @get()
+    def get_throttled(self):
+        return self.app_service.success()
+
+    @throttle(limit=5, ttl=10)
+    @get("/higher")
+    def get_higher(self):
+        return self.app_service.success()
+
+    @get("/shorter")
+    @throttle(limit=5, ttl=2)
+    def get_shorter(self):
+        return self.app_service.success()
+
+    @get("/shorter-2")
+    @throttle(limit=5, ttl=2)
+    def get_shorter_2(self):
+        return self.app_service.success()

tests/app/controller/default.py

@@ -0,0 +1,17 @@
+from ellar.common import Module
+
+from ellar_throttler import ThrottlerModule
+
+from .app import AppController
+from .default import DefaultController
+from .limit import LimitController
+from .service import AppService
+
+
+@Module(
+    modules=[ThrottlerModule.module_configure(limit=5, ttl=60)],
+    controllers=(DefaultController, LimitController, AppController),
+    providers=[AppService],
+)
+class ControllerModule:
+    pass

tests/app/controller/limit.py

@@ -0,0 +1,10 @@
+from ellar.di import injectable
+
+
+@injectable()
+class AppService:
+    def success(self):
+        return {"success": True}
+
+    def ignored(self):
+        return {"ignored": True}

tests/app/controller/module.py

@@ -0,0 +1,8 @@
+from ellar.common import Module
+
+from .controller.module import ControllerModule
+
+
+@Module(modules=(ControllerModule,))
+class AppModule:
+    pass