From b73e1088e01329684c0724081840af3870870f52 Mon Sep 17 00:00:00 2001
From: Lingxian Kong <anlin.kong@gmail.com>
Date: Wed, 13 Jan 2021 15:04:39 +1300
Subject: [PATCH] Restore test user's original roles during clean up

Depends-On: https://review.opendev.org/c/openstack/patrole/+/770867
Change-Id: I965d9e280ef0a85dfe6b2e4000eb47e1d06fed77
---
 patrole_tempest_plugin/rbac_utils.py           | 18 ++++++++++++++++++
 .../tests/unit/test_rbac_utils.py              | 14 ++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/patrole_tempest_plugin/rbac_utils.py b/patrole_tempest_plugin/rbac_utils.py
index b11ae383..7a1c33fb 100644
--- a/patrole_tempest_plugin/rbac_utils.py
+++ b/patrole_tempest_plugin/rbac_utils.py
@@ -143,9 +143,20 @@ class RbacUtilsMixin(object):
     _user_id = None
     _role_map = None
     _role_inferences_mapping = None
+    _orig_roles = []
 
     admin_roles_client = None
 
+    @classmethod
+    def restore_roles(cls):
+        if cls._orig_roles:
+            LOG.info("Restoring original roles %s", cls._orig_roles)
+            roles_already_present = cls._list_and_clear_user_roles_on_project(
+                cls._orig_roles)
+
+            if not roles_already_present:
+                cls._create_user_role_on_project(cls._orig_roles)
+
     @classmethod
     def setup_clients(cls):
         # Intialize the admin roles_client to perform role switching.
@@ -165,8 +176,15 @@ class RbacUtilsMixin(object):
 
         cls._init_roles()
 
+        # Store the user's original roles and rollback after testing.
+        roles = cls.admin_roles_client.list_user_roles_on_project(
+            cls._project_id, cls._user_id)['roles']
+        cls._orig_roles = [role['id'] for role in roles]
+        cls.addClassResourceCleanup(cls.restore_roles)
+
         # Change default role to admin
         cls._override_role(False)
+
         super(RbacUtilsMixin, cls).setup_clients()
 
     @classmethod
diff --git a/patrole_tempest_plugin/tests/unit/test_rbac_utils.py b/patrole_tempest_plugin/tests/unit/test_rbac_utils.py
index e61f798d..f72a4f4a 100644
--- a/patrole_tempest_plugin/tests/unit/test_rbac_utils.py
+++ b/patrole_tempest_plugin/tests/unit/test_rbac_utils.py
@@ -248,6 +248,20 @@ class RBACUtilsMixinTest(base.TestCase):
             actual_roles = sorted(self.test_obj.get_all_needed_roles(roles))
         self.assertEqual(expected_roles, actual_roles)
 
+    def test_restore_roles(self):
+        self.rbac_utils_fixture.set_roles(['admin', 'member'], 'member')
+        roles_client = self.rbac_utils_fixture.admin_roles_client
+
+        # Explicitly call setup_clients() to make sure cls._orig_roles is set
+        # properly. Explicitly call resource_cleanup to invoke restore_roles().
+        self.test_obj.setup_clients()
+        self.test_obj.resource_cleanup()
+
+        # list_user_roles_on_project is called twice in setup_clients(),
+        # restore_roles() is called twice during resource cleanup.
+        self.assertEqual(4, roles_client.list_user_roles_on_project.call_count)
+        self.assertEqual(['member_id'], self.test_obj._orig_roles)
+
 
 class ValidateListContextTest(base.TestCase):
     @staticmethod