Unit test for finding handlers

modal/__init__.py

@@ -19,8 +19,11 @@
     from .functions import (
         Function,
         asgi_app,
+        build,
         current_function_call_id,
         current_input_id,
+        enter,
+        exit,
         method,
         web_endpoint,
         wsgi_app,
@@ -66,10 +69,13 @@
     "Tunnel",
     "Volume",
     "asgi_app",
+    "build",
     "container_app",
     "create_package_mounts",
     "current_function_call_id",
     "current_input_id",
+    "enter",
+    "exit",
     "forward",
     "is_local",
     "method",

modal/functions.py

@@ -1714,6 +1714,8 @@ def wrapper(f: Union[Callable[[], Any], _PartialFunction]) -> _PartialFunction:
         else:
             return _PartialFunction(f, _PartialFunctionFlags.BUILD)
 
+    return wrapper
+
 
 @typechecked
 def _enter(_warn_parentheses_missing=None) -> Callable[[Union[Callable[[], Any], _PartialFunction]], _PartialFunction]:
@@ -1726,6 +1728,8 @@ def wrapper(f: Union[Callable[[], Any], _PartialFunction]) -> _PartialFunction:
         else:
             return _PartialFunction(f, _PartialFunctionFlags.ENTER)
 
+    return wrapper
+
 
 # TODO(erikbern): last argument should be Optional[TracebackType]
 ExitHandlerType = Callable[[Optional[Type[BaseException]], Optional[BaseException], Any], None]
@@ -1739,6 +1743,8 @@ def _exit(_warn_parentheses_missing=None) -> Callable[[ExitHandlerType], _Partia
     def wrapper(f: ExitHandlerType) -> _PartialFunction:
         return _PartialFunction(f, _PartialFunctionFlags.EXIT)
 
+    return wrapper
+
 
 method = synchronize_api(_method)
 web_endpoint = synchronize_api(_web_endpoint)

test/cls_test.py

@@ -6,11 +6,12 @@
 
 from typing_extensions import assert_type
 
-from modal import Cls, Function, Stub, method
+from modal import Cls, Dict, Function, Stub, build, enter, exit, method
 from modal._serialization import deserialize
 from modal.app import ContainerApp
 from modal.cls import ClsMixin
 from modal.exception import DeprecationError, ExecutionError
+from modal.functions import _find_partial_methods, _PartialFunction, _PartialFunctionFlags
 from modal.runner import deploy_stub
 from modal_proto import api_pb2
 from modal_test_support.base_class import BaseCls2
@@ -494,3 +495,35 @@ def test_keep_warm_depr():
 
     with pytest.warns(DeprecationError, match="@method"):
         stub.cls(keep_warm=2)(ClsWith2Methods)
+
+
+class ClsWithHandlers:
+    @build()
+    def my_build(self):
+        pass
+
+    @enter()
+    def my_enter(self):
+        pass
+
+    @build()
+    @enter()
+    def my_build_and_enter(self):
+        pass
+
+    @exit()
+    def my_exit(self, exc_type, exc, traceback):
+        pass
+
+
+def test_handlers():
+    pfs: Dict[str, _PartialFunction]
+
+    pfs = _find_partial_methods(ClsWithHandlers, _PartialFunctionFlags.BUILD)
+    assert list(pfs.keys()) == ["my_build", "my_build_and_enter"]
+
+    pfs = _find_partial_methods(ClsWithHandlers, _PartialFunctionFlags.ENTER)
+    assert list(pfs.keys()) == ["my_enter", "my_build_and_enter"]
+
+    pfs = _find_partial_methods(ClsWithHandlers, _PartialFunctionFlags.EXIT)
+    assert list(pfs.keys()) == ["my_exit"]