Add role to OrganizationMembership class (#236)

* Add role to OrganizationMembership class
* Add PUT OrganizationMembership endpoint

Author: ameesha

tests/test_user_management.py

@@ -319,6 +319,20 @@ def test_create_organization_membership(
         assert organization_membership["user_id"] == user_id
         assert organization_membership["organization_id"] == organization_id
 
+    def test_update_organization_membership(
+        self, capture_and_mock_request, mock_organization_membership
+    ):
+        url, _ = capture_and_mock_request("put", mock_organization_membership, 201)
+
+        organization_membership = self.user_management.update_organization_membership(
+            organization_membership_id="om_ABCDE",
+            role_slug="member",
+        )
+
+        assert url[0].endswith("user_management/organization_memberships/om_ABCDE")
+        assert organization_membership["id"] == "om_ABCDE"
+        assert organization_membership["role"] == {"slug": "member"}
+
     def test_get_organization_membership(
         self, mock_organization_membership, capture_and_mock_request
     ):

tests/utils/fixtures/mock_organization_membership.py

@@ -8,6 +8,7 @@ def __init__(self, id):
         self.user_id = "user_12345"
         self.organization_id = "org_67890"
         self.status = "active"
+        self.role = {"slug": "member"}
         self.created_at = datetime.datetime.now()
         self.updated_at = datetime.datetime.now()
 
@@ -16,6 +17,7 @@ def __init__(self, id):
         "user_id",
         "organization_id",
         "status",
+        "role",
         "created_at",
         "updated_at",
     ]

workos/__about__.py

@@ -12,7 +12,7 @@
 
 __package_url__ = "https://github.com/workos-inc/workos-python"
 
-__version__ = "3.0.0"
+__version__ = "3.1.0"
 
 __author__ = "WorkOS"
 

workos/resources/user_management.py

@@ -69,6 +69,7 @@ class WorkOSOrganizationMembership(WorkOSBaseResource):
         "status",
         "created_at",
         "updated_at",
+        "role",
     ]
 
 

workos/user_management.py

@@ -201,19 +201,25 @@ def delete_user(self, user_id):
             token=workos.api_key,
         )
 
-    def create_organization_membership(self, user_id, organization_id):
+    def create_organization_membership(self, user_id, organization_id, role_slug=None):
         """Create a new OrganizationMembership for the given Organization and User.
 
         Args:
             user_id: The Unique ID of the User.
             organization_id: The Unique ID of the Organization to which the user belongs to.
+            role_slug: The Unique Slug of the Role to which to grant to this membership.
+                If no slug is passed in, the default role will be granted.(Optional)
 
         Returns:
             dict: Created OrganizationMembership response from WorkOS.
         """
         headers = {}
 
-        params = {"user_id": user_id, "organization_id": organization_id}
+        params = {
+            "user_id": user_id,
+            "organization_id": organization_id,
+            "role_slug": role_slug,
+        }
 
         response = self.request_helper.request(
             ORGANIZATION_MEMBERSHIP_PATH,
@@ -225,6 +231,35 @@ def create_organization_membership(self, user_id, organization_id):
 
         return WorkOSOrganizationMembership.construct_from_response(response).to_dict()
 
+    def update_organization_membership(
+        self, organization_membership_id, role_slug=None
+    ):
+        """Updates an OrganizationMembership for the given id.
+
+        Args:
+            organization_membership_id (str) -  The unique ID of the Organization Membership.
+            role_slug: The Unique Slug of the Role to which to grant to this membership.
+                If no slug is passed in, it will not be changed (Optional)
+
+        Returns:
+            dict: Created OrganizationMembership response from WorkOS.
+        """
+        headers = {}
+
+        params = {
+            "role_slug": role_slug,
+        }
+
+        response = self.request_helper.request(
+            ORGANIZATION_MEMBERSHIP_DETAIL_PATH.format(organization_membership_id),
+            method=REQUEST_METHOD_PUT,
+            params=params,
+            headers=headers,
+            token=workos.api_key,
+        )
+
+        return WorkOSOrganizationMembership.construct_from_response(response).to_dict()
+
     def get_organization_membership(self, organization_membership_id):
         """Get the details of an organization membership.