Add test for removing env role in update_environment_roles

atst/domain/environments.py

@@ -58,10 +58,10 @@ class Environments(object):
         return env
 
     @classmethod
-    def update_environment_role(cls, user, ids_and_roles, workspace_user):
+    def update_environment_roles(cls, user, workspace, workspace_user, ids_and_roles):
         Authorization.check_workspace_permission(
             user,
-            workspace_user.workspace,
+            workspace,
             Permissions.ADD_AND_ASSIGN_CSP_ROLES,
             "assign environment roles",
         )
@@ -73,7 +73,9 @@ class Environments(object):
             if new_role is None:
                 EnvironmentRoles.delete(workspace_user.user.id, environment.id)
             else:
-                env_role = EnvironmentRoles.get(workspace_user.user_id, id_and_role["id"])
+                env_role = EnvironmentRoles.get(
+                    workspace_user.user.id, id_and_role["id"]
+                )
                 if env_role:
                     env_role.role = new_role
                 else:

atst/domain/projects.py

@@ -2,6 +2,7 @@ from atst.database import db
 from atst.domain.authz import Authorization
 from atst.domain.environments import Environments
 from atst.domain.exceptions import NotFoundError
+from atst.domain.environment_roles import EnvironmentRoles
 from atst.models.permissions import Permissions
 from atst.models.project import Project
 from atst.models.environment import Environment

atst/models/workspace_user.py

@@ -56,6 +56,18 @@ class WorkspaceUser(object):
             .count()
         )
 
+    @property
+    def environment_roles(self):
+        return (
+            db.session.query(EnvironmentRole)
+            .join(EnvironmentRole.environment)
+            .join(Environment.project)
+            .join(Project.workspace)
+            .filter(Project.workspace_id == self.workspace_id)
+            .filter(EnvironmentRole.user_id == self.user_id)
+            .all()
+        )
+
     @property
     def has_environment_roles(self):
         return self.num_environment_roles > 0

atst/routes/workspaces.py

@@ -297,7 +297,9 @@ def update_member(workspace_id, member_id):
             )
             new_role_name = member.role_displayname
 
-        Environments.update_environment_role(g.current_user, ids_and_roles, member)
+        Environments.update_environment_roles(
+            g.current_user, workspace, member, ids_and_roles
+        )
 
         return redirect(
             url_for(

tests/domain/test_environments.py

@@ -1,5 +1,6 @@
 from atst.domain.environments import Environments
 from atst.domain.environment_roles import EnvironmentRoles
+from atst.domain.workspace_users import WorkspaceUsers
 
 from tests.factories import UserFactory, WorkspaceFactory
 
@@ -37,7 +38,7 @@ def test_update_environment_roles():
     ]
 
     workspace_user = workspace.members[0]
-    Environments.update_environment_role(owner, new_ids_and_roles, workspace_user)
+    Environments.update_environment_roles(owner, workspace, workspace_user, new_ids_and_roles)
     new_dev_env_role = EnvironmentRoles.get(workspace_user.user.id, dev_env.id)
     staging_env_role = EnvironmentRoles.get(workspace_user.user.id, staging_env.id)
 
@@ -45,6 +46,57 @@ def test_update_environment_roles():
     assert staging_env_role.role == "developer"
 
 
+def test_remove_environment_role():
+    owner = UserFactory.create()
+    developer = UserFactory.from_atat_role("developer")
+    workspace = WorkspaceFactory.create(
+        owner=owner,
+        members=[{"user": developer, "role_name": "developer"}],
+        projects=[
+            {
+                "name": "project1",
+                "environments": [
+                    {
+                        "name": "project1 dev",
+                        "members": [{"user": developer, "role_name": "devops"}],
+                    },
+                    {
+                        "name": "project1 staging",
+                        "members": [{"user": developer, "role_name": "developer"}],
+                    },
+                    {
+                        "name": "project1 uat",
+                        "members": [
+                            {"user": developer, "role_name": "financial_auditor"}
+                        ],
+                    },
+                    {"name": "project1 prod"},
+                ],
+            }
+        ],
+    )
+
+    project = workspace.projects[0]
+    now_ba = project.environments[0].id
+    now_none = project.environments[1].id
+    still_fa = project.environments[2].id
+
+    new_environment_roles = [
+        {"id": now_ba, "role": "billing_auditor"},
+        {"id": now_none, "role": None},
+    ]
+
+    workspace_user = WorkspaceUsers.get(workspace.id, developer.id)
+    Environments.update_environment_roles(
+        owner, workspace, workspace_user, new_environment_roles
+    )
+
+    assert workspace_user.num_environment_roles == 2
+    assert EnvironmentRoles.get(developer.id, now_ba).role == "billing_auditor"
+    assert EnvironmentRoles.get(developer.id, now_none) is None
+    assert EnvironmentRoles.get(developer.id, still_fa).role == "financial_auditor"
+
+
 def test_get_scoped_environments(db):
     developer = UserFactory.create()
     workspace = WorkspaceFactory.create(