Merge pull request #15 from melvinkcx/bugfixes/local_handler_event_name

supporting Enum as event name

fastapi_events/dispatcher.py

@@ -1,12 +1,13 @@
 import os
 from distutils.util import strtobool
-from typing import Deque, Any, Optional
+from enum import Enum
+from typing import Deque, Any, Optional, Union
 
 from fastapi_events import event_store
 from fastapi_events.typing import Event
 
 
-def dispatch(event_name: str, payload: Optional[Any] = None) -> None:
+def dispatch(event_name: Union[str, Enum], payload: Optional[Any] = None) -> None:
     DISABLE_DISPATCH_GLOBALLY = strtobool(os.environ.get("FASTAPI_EVENTS_DISABLE_DISPATCH", "0"))
 
     if DISABLE_DISPATCH_GLOBALLY:

fastapi_events/handlers/local.py

@@ -36,12 +36,18 @@ async def handle(self, event: Event) -> None:
                 await loop.run_in_executor(None, functools.partial(handler, event))
 
     def _register_handler(self, event_name, func):
+        if not isinstance(event_name, str):
+            event_name = str(event_name)
+
         if event_name not in self._registry:
             self._registry[event_name] = []
 
         self._registry[event_name].append(func)
 
     def _get_handlers_for_event(self, event_name):
+        if not isinstance(event_name, str):
+            event_name = str(event_name)
+
         # TODO consider adding a cache
         handlers = []
         for event_name_pattern, registered_handlers in self._registry.items():

fastapi_events/typing.py

@@ -1,3 +1,4 @@
-from typing import Tuple, Any
+from enum import Enum
+from typing import Tuple, Any, Union
 
-Event = Tuple[str, Any]
+Event = Tuple[Union[str, Enum], Any]

tests/handlers/test_local_handler.py

@@ -1,3 +1,5 @@
+from enum import Enum
+
 from starlette.applications import Starlette
 from starlette.middleware import Middleware
 from starlette.requests import Request
@@ -66,3 +68,32 @@ async def root(request: Request) -> JSONResponse:
         ("dance", 1)
     ):
         assert len(events_handled[event_category]) == expected_count
+
+
+def test_local_handler_with_enum():
+    """
+    Test local_handler with Enum as event name
+    """
+
+    class Events(Enum):
+        CREATED = "CREATED"
+
+    events_handled = []
+
+    @local_handler.register(event_name=Events.CREATED)
+    async def handle_all_created_events(event: Event):
+        events_handled.append(event)
+
+    app = Starlette(middleware=[
+        Middleware(EventHandlerASGIMiddleware,
+                   handlers=[local_handler])])
+
+    @app.route("/")
+    async def root(request: Request) -> JSONResponse:
+        dispatch(Events.CREATED)
+        return JSONResponse()
+
+    client = TestClient(app)
+    client.get("/")
+
+    assert events_handled[0][0] == Events.CREATED