fix(test): make test_clean_requests_after_schema_grant more idempotent

apache · Apr 8, 2022 · 00137bc · 00137bc
1 parent d52e386
commit 00137bc
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/tests/integration_tests/access_tests.py b/tests/integration_tests/access_tests.py
@@ -374,6 +374,7 @@ def test_clean_requests_after_schema_grant(self):
             .filter_by(table_name="wb_health_population")
             .first()
         )
+        original_schema = ds.schema
 
         ds.schema = "temp_schema"
         security_manager.add_permission_view_menu("schema_access", ds.schema_perm)
@@ -394,13 +395,7 @@ def test_clean_requests_after_schema_grant(self):
         gamma_user = security_manager.find_user(username="gamma")
         gamma_user.roles.remove(security_manager.find_role(SCHEMA_ACCESS_ROLE))
 
-        ds = (
-            session.query(SqlaTable)
-            .filter_by(table_name="wb_health_population")
-            .first()
-        )
-        ds.schema = None
-
+        ds.schema = original_schema
         session.commit()
 
     @mock.patch("superset.utils.core.send_mime_email")