Test passing

resonatehq · Jan 27, 2025 · afb8f3c · afb8f3c
1 parent eb2f272
commit afb8f3c
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 10 deletions.
diff --git a/src/resonate/resonate.py b/src/resonate/resonate.py
@@ -181,15 +181,28 @@ def register(
             (func, Options(version=version, durable=True, retry_policy=retry_policy)),
         )
 
-    def fn(
+    def __call__(
         self,
         name: str | None = None,
         version: int = 1,
         retry_policy: retry_policy.RetryPolicy | None = None,
     ) -> Callable[
         [Callable[Concatenate[Context, P], Any]],
         RegisteredFn[P, Any],
-    ]: ...
+    ]:
+        def wrapper(
+            func: Callable[Concatenate[Context, P], Any],
+        ) -> RegisteredFn[P, Any]:
+            self._registry.add(
+                name or func.__name__,
+                (
+                    func,
+                    Options(version=version, durable=True, retry_policy=retry_policy),
+                ),
+            )
+            return RegisteredFn(self._scheduler, func)
+
+        return wrapper
 
     @overload
     def run(

diff --git a/tests/test_functionality.py b/tests/test_functionality.py
@@ -459,7 +459,7 @@ def test_golden_device_rfi_and_lfc_with_decorator() -> None:
         task_source=Poller("http://localhost:8002", group=group),
     )
 
-    @resonate.fn()
+    @resonate()
     def foo(ctx: Context, n: str) -> Generator[Yieldable, Any, str]:
         v: str = yield ctx.lfc(bar, n).options(
             id="bar",
@@ -472,7 +472,7 @@ def bar(ctx: Context, n: str) -> Generator[Yieldable, Any, str]:
         v: str = yield p
         return v
 
-    @resonate.fn()
+    @resonate()
     def baz(ctx: Context, n: str) -> str:  # noqa: ARG001
         return n
 
@@ -611,7 +611,7 @@ def test_golden_device_rfc_and_lfc_with_decorator() -> None:
         task_source=Poller("http://localhost:8002", group=group),
     )
 
-    @resonate.fn()
+    @resonate()
     def foo(ctx: Context, n: str) -> Generator[Yieldable, Any, str]:
         v: str = yield ctx.lfc(bar, n).options(
             id="bar",
@@ -623,7 +623,7 @@ def bar(ctx: Context, n: str) -> Generator[Yieldable, Any, str]:
         v: str = yield ctx.rfc(baz, n).options(id="baz", send_to=poll(group))
         return v
 
-    @resonate.fn()
+    @resonate()
     def baz(ctx: Context, n: str) -> str:  # noqa: ARG001
         return n
 
@@ -734,7 +734,7 @@ def test_sleep() -> None:
         store=store, task_source=Poller("http://localhost:8002", group=group)
     )
 
-    @resonate.fn()
+    @resonate()
     def foo_sleep(ctx: Context, n: int) -> Generator[Yieldable, Any, int]:
         yield ctx.sleep(n)
         return n
@@ -794,7 +794,7 @@ def test_golden_device_detached_with_registered() -> None:
         task_source=Poller("http://localhost:8002", group=group),
     )
 
-    @resonate.fn()
+    @resonate()
     def foo_golden_device_detached_with_registered(
         ctx: Context, n: str
     ) -> Generator[Yieldable, Any, str]:
@@ -807,11 +807,11 @@ def foo_golden_device_detached_with_registered(
         v: str = yield p
         return v
 
-    @resonate.fn()
+    @resonate()
     def bar_golden_device_detached_with_registered(ctx: Context, n: str) -> str:  # noqa: ARG001
         return n
 
-    @resonate.fn()
+    @resonate()
     def baz_golden_device_detached_with_registered(
         ctx: Context,  # noqa: ARG001
         promise_id: str,