Apply suggested changes

workos · Nov 20, 2023 · 84f68eb · 84f68eb
1 parent e0bbeaa
commit 84f68eb
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 62 deletions.
diff --git a/tests/test_user_management.py b/tests/test_user_management.py
@@ -128,14 +128,6 @@ def mock_magic_auth_challenge_response(self):
     @pytest.fixture
     def mock_enroll_auth_factor_response(self):
         return {
-            "authentication_challenge": {
-                "object": "authentication_challenge",
-                "id": "auth_challenge_01FVYZWQTZQ5VB6BC5MPG2EYC5",
-                "created_at": "2022-02-15T15:26:53.274Z",
-                "updated_at": "2022-02-15T15:26:53.274Z",
-                "expires_at": "2022-02-15T15:36:53.279Z",
-                "authentication_factor_id": "auth_factor_01FVYZ5QM8N98T9ME5BCB2BBMJ",
-            },
             "authentication_factor": {
                 "object": "authentication_factor",
                 "id": "auth_factor_01FVYZ5QM8N98T9ME5BCB2BBMJ",
@@ -148,6 +140,14 @@ def mock_enroll_auth_factor_response(self):
                     "uri": "otpauth://totp/FooCorp:[email protected]?secret=NAGCCFS3EYRB422HNAKAKY3XDUORMSRF&issuer=FooCorp",
                 },
             },
+            "authentication_challenge": {
+                "object": "authentication_challenge",
+                "id": "auth_challenge_01FVYZWQTZQ5VB6BC5MPG2EYC5",
+                "created_at": "2022-02-15T15:26:53.274Z",
+                "updated_at": "2022-02-15T15:26:53.274Z",
+                "expires_at": "2022-02-15T15:36:53.279Z",
+                "authentication_factor_id": "auth_factor_01FVYZ5QM8N98T9ME5BCB2BBMJ",
+            },
         }
 
     @pytest.fixture

diff --git a/workos/resources/mfa.py b/workos/resources/mfa.py
@@ -115,51 +115,3 @@ def to_dict(self):
         verification_response_dict = super(WorkOSChallengeVerification, self).to_dict()
 
         return verification_response_dict
-
-
-class WorkOSAuthenticationChallengeAndFactor(WorkOSBaseResource):
-    """Representation of an Authentication Challenge and Factor as returned by WorkOS through the User Management feature.
-
-    Attributes:
-        OBJECT_FIELDS (list): List of fields a WorkOSAuthenticationChallengeAndFactor is comprised of.
-    """
-
-    OBJECT_FIELDS = [
-        "authentication_challenge",
-        "authentication_factor",
-    ]
-
-    @classmethod
-    def construct_from_response(cls, response):
-        authentication_challenge_and_factor = super(
-            WorkOSAuthenticationChallengeAndFactor, cls
-        ).construct_from_response(response)
-
-        authentication_challenge_and_factor.authentication_challenge = (
-            WorkOSChallenge.construct_from_response(
-                response["authentication_challenge"]
-            )
-        )
-
-        authentication_challenge_and_factor.authentication_factor = (
-            WorkOSAuthenticationFactorTotp.construct_from_response(
-                response["authentication_factor"]
-            )
-        )
-
-        return authentication_challenge_and_factor
-
-    def to_dict(self):
-        authentication_challenge_and_factor_dict = super(
-            WorkOSAuthenticationChallengeAndFactor, self
-        ).to_dict()
-
-        challenge_dict = self.authentication_challenge.to_dict()
-        authentication_challenge_and_factor_dict[
-            "authentication_challenge"
-        ] = challenge_dict
-
-        factor_dict = self.authentication_factor.to_dict()
-        authentication_challenge_and_factor_dict["authentication_factor"] = factor_dict
-
-        return authentication_challenge_and_factor_dict
diff --git a/workos/user_management.py b/workos/user_management.py
@@ -2,8 +2,7 @@
 from workos.resources.authentication_response import WorkOSAuthenticationResponse
 from workos.resources.password_challenge_response import WorkOSPasswordChallengeResponse
 from workos.resources.list import WorkOSListResource
-from workos.resources.mfa import WorkOSAuthenticationChallengeAndFactor
-from workos.resources.mfa import WorkOSAuthenticationFactorTotp
+from workos.resources.mfa import WorkOSAuthenticationFactorTotp, WorkOSChallenge
 from workos.resources.users import WorkOSUser
 from workos.utils.pagination_order import Order
 from workos.utils.request import (
@@ -557,8 +556,7 @@ def enroll_auth_factor(
             totp_issuer (str): Name of the Organization (Optional)
             totp_user (str): Email of user (Optional)
 
-        Returns:
-            dict: AuthenticationChallengeAndFactor response from WorkOS.
+        Returns: { WorkOSAuthenticationFactorTotp, WorkOSChallenge}
         """
 
         if type not in ["totp"]:
@@ -580,9 +578,19 @@ def enroll_auth_factor(
             token=workos.api_key,
         )
 
-        return WorkOSAuthenticationChallengeAndFactor.construct_from_response(
-            response
+        factor_and_challenge = {}
+        factor_and_challenge[
+            "authentication_factor"
+        ] = WorkOSAuthenticationFactorTotp.construct_from_response(
+            response["authentication_factor"]
         ).to_dict()
+        factor_and_challenge[
+            "authentication_challenge"
+        ] = WorkOSChallenge.construct_from_response(
+            response["authentication_challenge"]
+        ).to_dict()
+
+        return factor_and_challenge
 
     def list_auth_factors(
         self,