feat: Add tests for email and sms blocking triggers, format code, solve cyclic import issue with event types

exaby73 · exaby73 · commit 6ba3a046d1cf · 2024-12-05T15:43:25.000+05:30
diff --git a/samples/identity/functions/main.py b/samples/identity/functions/main.py
@@ -11,7 +11,7 @@ def beforeusercreated(
     event: identity_fn.AuthBlockingEvent
 ) -> identity_fn.BeforeCreateResponse | None:
     print(event)
-    if not event.data.email:
+    if not event.data or not event.data.email:
         return None
     if "@cats.com" in event.data.email:
         return identity_fn.BeforeCreateResponse(display_name="Meow!",)
@@ -29,7 +29,7 @@ def beforeusersignedin(
     event: identity_fn.AuthBlockingEvent
 ) -> identity_fn.BeforeSignInResponse | None:
     print(event)
-    if not event.data.email:
+    if not event.data or not event.data.email:
         return None
 
     if "@cats.com" in event.data.email:
diff --git a/src/firebase_functions/firestore_fn.py b/src/firebase_functions/firestore_fn.py
@@ -218,8 +218,11 @@ def _firestore_endpoint_handler(
                                                      auth_id=event_auth_id)
         func(database_event_with_auth_context)
     else:
-        # mypy cannot infer that the event type is correct, hence the cast
-        _typing.cast(_C1 | _C2, func)(database_event)
+        # Split the casting into two separate branches based on event type
+        if event_type in (_event_type_written, _event_type_updated):
+            _typing.cast(_C1, func)(_typing.cast(_E1, database_event))
+        else:
+            _typing.cast(_C2, func)(_typing.cast(_E2, database_event))
 
 
 @_util.copy_func_kwargs(FirestoreOptions)
diff --git a/src/firebase_functions/identity_fn.py b/src/firebase_functions/identity_fn.py
@@ -243,16 +243,19 @@ class Credential:
     sign_in_method: str
     """The user's sign-in method."""
 
+
 EmailType = _typing.Literal["EMAIL_SIGN_IN", "PASSWORD_RESET"]
-SmsType = _typing.Literal["SIGN_IN_OR_SIGN_UP", "MULTI_FACTOR_SIGN_IN", "MULTI_FACTOR_ENROLLMENT"]
+SmsType = _typing.Literal["SIGN_IN_OR_SIGN_UP", "MULTI_FACTOR_SIGN_IN",
+                          "MULTI_FACTOR_ENROLLMENT"]
+
 
 @_dataclasses.dataclass(frozen=True)
 class AuthBlockingEvent:
     """
     Defines an auth event for identitytoolkit v2 auth blocking events.
     """
 
-    data: AuthUserRecord | None # This is None for beforeEmailSent and beforeSmsSent events
+    data: AuthUserRecord | None  # This is None for beforeEmailSent and beforeSmsSent events
     """
     The UserRecord passed to auth blocking functions from the identity platform.
     """
@@ -345,7 +348,7 @@ class BeforeEmailSentResponse(_typing.TypedDict, total=False):
     recaptcha_action_override: RecaptchaActionOptions | None
 
 
-class BeforeSmsSentResponse(BeforeEmailSentResponse, total=False):
+class BeforeSmsSentResponse(_typing.TypedDict, total=False):
     """
     The handler response type for 'before_sms_sent' blocking events.
     """
@@ -372,7 +375,7 @@ class BeforeSmsSentResponse(BeforeEmailSentResponse, total=False):
 """
 
 BeforeSmsSentCallable = _typing.Callable[[AuthBlockingEvent],
-                                          BeforeSmsSentResponse | None]
+                                         BeforeSmsSentResponse | None]
 """
 The type of the callable for 'before_sms_sent' blocking events.
 """
@@ -485,6 +488,7 @@ def before_user_created_wrapped(request: _Request) -> _Response:
 
     return before_user_created_decorator
 
+
 @_util.copy_func_kwargs(_options.BlockingOptions)
 def before_email_sent(
     **kwargs,
@@ -499,20 +503,23 @@ def before_email_sent(
       from firebase_functions import identity_fn
 
       @identity_fn.before_email_sent()
-      def example(event: identity_fn.AuthBlockingEvent) -> identity_fn.BeforeEmailSentResponse | None:
+      def example(
+          event: identity_fn.AuthBlockingEvent
+      ) -> identity_fn.BeforeEmailSentResponse | None:
           pass
 
     :param \\*\\*kwargs: Options.
     :type \\*\\*kwargs: as :exc:`firebase_functions.options.BlockingOptions`
     :rtype: :exc:`typing.Callable`
             \\[ \\[ :exc:`firebase_functions.identity_fn.AuthBlockingEvent` \\],
             :exc:`firebase_functions.identity_fn.BeforeEmailSentResponse` \\| `None` \\]
-            A function that takes a AuthBlockingEvent and optionally returns BeforeEmailSentResponse.
+            A function that takes a AuthBlockingEvent and optionally returns 
+            BeforeEmailSentResponse.
     """
     options = _options.BlockingOptions(**kwargs)
 
     def before_email_sent_decorator(func: BeforeEmailSentCallable):
-        from firebase_functions.private._identity_fn import event_type_before_email_sent
+        from firebase_functions.private._identity_fn_event_types import event_type_before_email_sent
 
         @_functools.wraps(func)
         def before_email_sent_wrapped(request: _Request) -> _Response:
@@ -566,7 +573,7 @@ def example(event: identity_fn.AuthBlockingEvent) -> identity_fn.BeforeSmsSentRe
     options = _options.BlockingOptions(**kwargs)
 
     def before_sms_sent_decorator(func: BeforeSmsSentCallable):
-        from firebase_functions.private._identity_fn import event_type_before_sms_sent
+        from firebase_functions.private._identity_fn_event_types import event_type_before_sms_sent
 
         @_functools.wraps(func)
         def before_sms_sent_wrapped(request: _Request) -> _Response:
@@ -590,4 +597,4 @@ def before_sms_sent_wrapped(request: _Request) -> _Response:
         )
         return before_sms_sent_wrapped
 
-    return before_sms_sent_decorator
+    return before_sms_sent_decorator
diff --git a/src/firebase_functions/options.py b/src/firebase_functions/options.py
@@ -994,17 +994,20 @@ def _endpoint(
         self,
         **kwargs,
     ) -> _manifest.ManifestEndpoint:
-        from firebase_functions.private._identity_fn import event_type_before_create, event_type_before_sign_in
+        from firebase_functions.private._identity_fn_event_types import event_type_before_create, event_type_before_sign_in
 
         assert kwargs["event_type"] is not None
 
         blocking_trigger_options: _manifest.BlockingTriggerOptions
-        
-        if kwargs["event_type"] == event_type_before_create or kwargs["event_type"] == event_type_before_sign_in:
+
+        if kwargs["event_type"] == event_type_before_create or kwargs[
+                "event_type"] == event_type_before_sign_in:
             blocking_trigger_options = _manifest.BlockingTriggerOptions(
                 idToken=self.id_token if self.id_token is not None else False,
-                accessToken=self.access_token if self.access_token is not None else False,
-                refreshToken=self.refresh_token if self.refresh_token is not None else False,
+                accessToken=self.access_token
+                if self.access_token is not None else False,
+                refreshToken=self.refresh_token
+                if self.refresh_token is not None else False,
             )
         else:
             blocking_trigger_options = _manifest.BlockingTriggerOptions()
diff --git a/src/firebase_functions/private/_identity_fn.py b/src/firebase_functions/private/_identity_fn.py
@@ -20,6 +20,7 @@
 
 from firebase_functions.core import _with_init
 from firebase_functions.https_fn import HttpsError, FunctionsErrorCode
+from firebase_functions.private._identity_fn_event_types import event_type_before_create, event_type_before_sign_in
 
 import firebase_functions.private.util as _util
 import firebase_functions.private.token_verifier as _token_verifier
@@ -202,11 +203,12 @@ def _credential_from_token_data(token_data: dict[str, _typing.Any],
     )
 
 
-def _auth_blocking_event_from_token_data(token_data: dict[str, _typing.Any], event_type: str):
+def _auth_blocking_event_from_token_data(token_data: dict[str, _typing.Any],
+                                         event_type: str):
     from firebase_functions.identity_fn import AuthBlockingEvent, AuthUserRecord
 
     data: AuthUserRecord | None = None
-    if event_type == event_type_before_create or event_type == event_type_before_sign_in:
+    if event_type in (event_type_before_create, event_type_before_sign_in):
         data = _auth_user_record_from_token_data(token_data["user_record"])
 
     return AuthBlockingEvent(
@@ -223,12 +225,6 @@ def _auth_blocking_event_from_token_data(token_data: dict[str, _typing.Any], eve
     )
 
 
-event_type_before_create = "providers/cloud.auth/eventTypes/user.beforeCreate"
-event_type_before_sign_in = "providers/cloud.auth/eventTypes/user.beforeSignIn"
-event_type_before_email_sent = "providers/cloud.auth/eventTypes/user.beforeSendEmail"
-event_type_before_sms_sent = "providers/cloud.auth/eventTypes/user.beforeSendSms"
-
-
 def _validate_auth_response(
     event_type: str,
     auth_response,
@@ -363,8 +359,8 @@ def before_operation_handler(
         jwt_token = request.json["data"]["jwt"]
         decoded_token = _token_verifier.verify_auth_blocking_token(jwt_token)
         event = _auth_blocking_event_from_token_data(decoded_token, event_type)
-        auth_response: BeforeCreateResponse | BeforeSignInResponse | BeforeEmailSentResponse | BeforeSmsSentResponse | None = _with_init(
-            func)(event)
+        auth_response: BeforeCreateResponse | BeforeSignInResponse | BeforeEmailSentResponse | \
+            BeforeSmsSentResponse | None = _with_init(func)(event)
         if not auth_response:
             return _jsonify({})
         auth_response_dict = _validate_auth_response(event_type, auth_response)
diff --git a/src/firebase_functions/private/_identity_fn_event_types.py b/src/firebase_functions/private/_identity_fn_event_types.py
@@ -0,0 +1,10 @@
+"""
+Identity function event types.
+"""
+
+# We need to import these from the identity_fn module, but due to circular import
+# issues, we need to define them here.
+event_type_before_create = "providers/cloud.auth/eventTypes/user.beforeCreate"
+event_type_before_sign_in = "providers/cloud.auth/eventTypes/user.beforeSignIn"
+event_type_before_email_sent = "providers/cloud.auth/eventTypes/user.beforeSendEmail"
+event_type_before_sms_sent = "providers/cloud.auth/eventTypes/user.beforeSendSms"
diff --git a/tests/test_identity_fn.py b/tests/test_identity_fn.py
@@ -7,8 +7,6 @@
 from flask import Flask, Request
 from werkzeug.test import EnvironBuilder
 
-from firebase_functions import core, identity_fn
-
 token_verifier_mock = MagicMock()
 token_verifier_mock.verify_auth_blocking_token = Mock(
     return_value={
@@ -24,8 +22,14 @@
         "user_agent": "user_agent",
         "iat": 0
     })
+
+firebase_admin_mock = MagicMock()
+firebase_admin_mock.initialize_app = Mock()
+firebase_admin_mock.get_app = Mock()
+
 mocked_modules = {
     "firebase_functions.private.token_verifier": token_verifier_mock,
+    "firebase_admin": firebase_admin_mock
 }
 
 
@@ -37,12 +41,14 @@ class TestIdentity(unittest.TestCase):
     def test_calls_init_function(self):
         hello = None
 
-        @core.init
-        def init():
-            nonlocal hello
-            hello = "world"
-
         with patch.dict("sys.modules", mocked_modules):
+            from firebase_functions import core, identity_fn
+
+            @core.init
+            def init():
+                nonlocal hello
+                hello = "world"
+
             app = Flask(__name__)
 
             func = Mock(__name__="example_func",
@@ -62,3 +68,84 @@ def init():
                 decorated_func(request)
 
         self.assertEqual("world", hello)
+
+    def test_auth_blocking_event_from_token_data_email(self):
+        """Test parsing a beforeSendEmail event."""
+        # Mock token data for email event
+        token_data = {
+            "iss": "https://securetoken.google.com/project_id",
+            "aud": "https://us-east1-project_id.cloudfunctions.net/function-1",
+            "iat": 1,  # Unix timestamp
+            "exp": 60 * 60 + 1,
+            "event_id": "EVENT_ID",
+            "event_type": "beforeSendEmail",
+            "user_agent": "USER_AGENT",
+            "ip_address": "1.2.3.4",
+            "locale": "en",
+            "recaptcha_score": 0.9,
+            "email_type": "PASSWORD_RESET",
+            "email": "johndoe@gmail.com"
+        }
+
+        with patch.dict("sys.modules", mocked_modules):
+            from firebase_functions.private._identity_fn import _auth_blocking_event_from_token_data
+            from firebase_functions.private._identity_fn_event_types import event_type_before_email_sent
+            import datetime
+
+            event = _auth_blocking_event_from_token_data(
+                token_data, event_type_before_email_sent)
+
+        self.assertEqual(event.event_id, "EVENT_ID")
+        self.assertEqual(event.ip_address, "1.2.3.4")
+        self.assertEqual(event.user_agent, "USER_AGENT")
+        self.assertEqual(event.locale, "en")
+        self.assertEqual(event.email_type, "PASSWORD_RESET")
+        self.assertEqual(event.sms_type, None)
+        self.assertEqual(event.data, None)  # No user record for email events
+        self.assertEqual(event.timestamp, datetime.datetime.fromtimestamp(1))
+
+        self.assertEqual(event.additional_user_info.email, "johndoe@gmail.com")
+        self.assertEqual(event.additional_user_info.recaptcha_score, 0.9)
+        self.assertEqual(event.additional_user_info.is_new_user, False)
+        self.assertEqual(event.additional_user_info.phone_number, None)
+
+    def test_auth_blocking_event_from_token_data_sms(self):
+        """Test parsing a beforeSendSms event."""
+        import datetime
+
+        token_data = {
+            "iss": "https://securetoken.google.com/project_id",
+            "aud": "https://us-east1-project_id.cloudfunctions.net/function-1",
+            "iat": 1,  # Unix timestamp
+            "exp": 60 * 60 + 1,
+            "event_id": "EVENT_ID",
+            "event_type": "beforeSendSms",
+            "user_agent": "USER_AGENT",
+            "ip_address": "1.2.3.4",
+            "locale": "en",
+            "recaptcha_score": 0.9,
+            "sms_type": "SIGN_IN_OR_SIGN_UP",
+            "phone_number": "+11234567890"
+        }
+
+        with patch.dict("sys.modules", mocked_modules):
+            from firebase_functions.private._identity_fn import _auth_blocking_event_from_token_data
+            from firebase_functions.private._identity_fn_event_types import event_type_before_sms_sent
+
+            event = _auth_blocking_event_from_token_data(
+                token_data, event_type_before_sms_sent)
+
+        self.assertEqual(event.event_id, "EVENT_ID")
+        self.assertEqual(event.ip_address, "1.2.3.4")
+        self.assertEqual(event.user_agent, "USER_AGENT")
+        self.assertEqual(event.locale, "en")
+        self.assertEqual(event.email_type, None)
+        self.assertEqual(event.sms_type, "SIGN_IN_OR_SIGN_UP")
+        self.assertEqual(event.data, None)  # No user record for SMS events
+        self.assertEqual(event.timestamp, datetime.datetime.fromtimestamp(1))
+
+        self.assertEqual(event.additional_user_info.phone_number,
+                         "+11234567890")
+        self.assertEqual(event.additional_user_info.recaptcha_score, 0.9)
+        self.assertEqual(event.additional_user_info.is_new_user, False)
+        self.assertEqual(event.additional_user_info.email, None)
