Repository: skypilot-org/skypilot
Pull request #5942 (2025-06-10 15:27:19)

lazy initialization of permissions service
<!-- Describe the changes in this PR -->
Lazily initialize the permissions service on first request. This has the effect of saving the initialization on the client side.


<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: install pre-commit (auto-check on commit) or `bash format.sh`
- [x] CI tests
- [x] All smoke tests: `/smoke-test` (CI) or `pytest tests/test_smoke.py` (local)
- [x] Backward compatibility: `/quicktest-core` (CI) or `pytest tests/smoke_tests/test_backward_compat.py` (local)
- [x] Manual tests: when using postgres DB backend, delete `state.db` and restart the API service. Verify that `state.db` file is not created on client side.
<!-- CI commands (/-prefixed) can only be triggered by repo members -->


SeungjinYang: /quicktest-core
SeungjinYang: /smoke-test -k test_workspace_switching

Patch:
diff --git a/sky/users/permission.py b/sky/users/permission.py
index 3252c1d05..13f41bcd5 100644
--- a/sky/users/permission.py
+++ b/sky/users/permission.py
@@ -30,26 +30,34 @@ class PermissionService:
     """Permission service for SkyPilot API Server."""
 
     def __init__(self):
-        global _enforcer_instance
-        if _enforcer_instance is None:
-            # For different threads, we share the same enforcer instance.
-            with _lock:
-                if _enforcer_instance is None:
-                    _enforcer_instance = self
-                    engine = global_user_state.initialize_and_get_db()
-                    adapter = sqlalchemy_adapter.Adapter(engine)
-                    model_path = os.path.join(os.path.dirname(__file__),
-                                              'model.conf')
-                    enforcer = casbin.Enforcer(model_path, adapter)
-                    self.enforcer = enforcer
-        else:
-            self.enforcer = _enforcer_instance.enforcer
-        with _policy_lock():
-            self._maybe_initialize_policies()
+        self.enforcer = None
+        self.init_lock = threading.Lock()
+
+    def _lazy_initialize(self):
+        if self.enforcer is not None:
+            return
+        with self.init_lock:
+            if self.enforcer is not None:
+                return
+            global _enforcer_instance
+            if _enforcer_instance is None:
+                # For different threads, we share the same enforcer instance.
+                with _lock:
+                    if _enforcer_instance is None:
+                        _enforcer_instance = self
+                        engine = global_user_state.initialize_and_get_db()
+                        adapter = sqlalchemy_adapter.Adapter(engine)
+                        model_path = os.path.join(os.path.dirname(__file__),
+                                                  'model.conf')
+                        enforcer = casbin.Enforcer(model_path, adapter)
+                        self.enforcer = enforcer
+            else:
+                self.enforcer = _enforcer_instance.enforcer
+            with _policy_lock():
+                self._maybe_initialize_policies()
 
     def _maybe_initialize_policies(self) -> None:
         """Initialize policies if they don't already exist."""
-        # TODO(zhwu): we should avoid running this on client side.
         logger.debug(f'Initializing policies in process: {os.getpid()}')
         self._load_policy_no_lock()
 
@@ -128,6 +136,7 @@ class PermissionService:
 
     def add_user_if_not_exists(self, user_id: str) -> None:
         """Add user role relationship."""
+        self._lazy_initialize()
         with _policy_lock():
             self._add_user_if_not_exists_no_lock(user_id)
 
@@ -147,6 +156,7 @@ class PermissionService:
 
     def update_role(self, user_id: str, new_role: str) -> None:
         """Update user role relationship."""
+        self._lazy_initialize()
         with _policy_lock():
             # Get current roles
             self._load_policy_no_lock()
@@ -179,6 +189,7 @@ class PermissionService:
         Returns:
             A list of role names that the user has.
         """
+        self._lazy_initialize()
         self._load_policy_no_lock()
         return self.enforcer.get_roles_for_user(user_id)
 
@@ -191,6 +202,7 @@ class PermissionService:
         # it is a hot path in every request. It is ok to have a stale policy,
         # as long as it is eventually consistent.
         # self._load_policy_no_lock()
+        self._lazy_initialize()
         return self.enforcer.enforce(user_id, path, method)
 
     def _load_policy_no_lock(self):
@@ -199,6 +211,7 @@ class PermissionService:
 
     def load_policy(self):
         """Load policy from storage with lock."""
+        self._lazy_initialize()
         with _policy_lock():
             self._load_policy_no_lock()
 
@@ -214,6 +227,7 @@ class PermissionService:
         For public workspaces, the permission is granted via a wildcard policy
         ('*').
         """
+        self._lazy_initialize()
         if os.getenv(constants.ENV_VAR_IS_SKYPILOT_SERVER) is None:
             # When it is not on API server, we allow all users to access all
             # workspaces, as the workspace check has been done on API server.
@@ -241,6 +255,7 @@ class PermissionService:
                    For public workspaces, this should be ['*'].
                    For private workspaces, this should be specific user IDs.
         """
+        self._lazy_initialize()
         with _policy_lock():
             for user in users:
                 logger.debug(f'Adding workspace policy: user={user}, '
@@ -258,6 +273,7 @@ class PermissionService:
                    For public workspaces, this should be ['*'].
                    For private workspaces, this should be specific user IDs.
         """
+        self._lazy_initialize()
         with _policy_lock():
             self._load_policy_no_lock()
             # Remove all existing policies for this workspace
@@ -271,6 +287,7 @@ class PermissionService:
 
     def remove_workspace_policy(self, workspace_name: str) -> None:
         """Remove workspace policy."""
+        self._lazy_initialize()
         with _policy_lock():
             self.enforcer.remove_filtered_policy(1, workspace_name)
             self.enforcer.save_policy()


Test patch:
diff --git a/tests/kubernetes/scripts/helm_gcp.sh b/tests/kubernetes/scripts/helm_gcp.sh
index d68f74cbe..ea387ebbc 100644
--- a/tests/kubernetes/scripts/helm_gcp.sh
+++ b/tests/kubernetes/scripts/helm_gcp.sh
@@ -98,10 +98,35 @@ ENDPOINT=http://${WEB_USERNAME}:${WEB_PASSWORD}@${HOST}
 
 echo "API server endpoint: $ENDPOINT"
 
-# Test the API server
+# Test the API server with retry logic
 echo "Testing API server health endpoint..."
-HEALTH_RESPONSE=$(curl -s ${ENDPOINT}/api/health)
-echo "Health response: $HEALTH_RESPONSE"
+echo "Endpoint: $ENDPOINT/api/health"
+MAX_RETRIES=10  # 5 minutes with 30 second intervals
+RETRY_INTERVAL=30
+RETRY_COUNT=0
+
+while [ $RETRY_COUNT -lt $MAX_RETRIES ]; do
+    echo "Attempt $((RETRY_COUNT + 1))/$MAX_RETRIES: Testing health endpoint..."
+    HEALTH_RESPONSE=$(curl -s ${ENDPOINT}/api/health)
+    echo "Health response: $HEALTH_RESPONSE"
+
+    # Check if response is valid JSON (not HTML error page)
+    if echo "$HEALTH_RESPONSE" | jq . >/dev/null 2>&1; then
+        echo "API server is responding with valid JSON!"
+        break
+    else
+        echo "API server not ready yet (received HTML error page). Waiting ${RETRY_INTERVAL} seconds..."
+        if [ $RETRY_COUNT -lt $((MAX_RETRIES - 1)) ]; then
+            sleep $RETRY_INTERVAL
+        fi
+        RETRY_COUNT=$((RETRY_COUNT + 1))
+    fi
+done
+
+if [ $RETRY_COUNT -eq $MAX_RETRIES ]; then
+    echo "Error: API server failed to become ready after $((MAX_RETRIES * RETRY_INTERVAL)) seconds"
+    exit 1
+fi
 
 # Extract version from response and verify it matches
 RETURNED_VERSION=$(echo $HEALTH_RESPONSE | jq -r '.version')
diff --git a/tests/unit_tests/test_sky/users/test_permission.py b/tests/unit_tests/test_sky/users/test_permission.py
index 43c5006dc..2a3f6aaf8 100644
--- a/tests/unit_tests/test_sky/users/test_permission.py
+++ b/tests/unit_tests/test_sky/users/test_permission.py
@@ -65,6 +65,7 @@ class TestPermissionService:
         with mock.patch.object(permission.PermissionService,
                                '_maybe_initialize_policies'):
             service = permission.PermissionService()
+            service._lazy_initialize()
 
             # Verify SQLAlchemy adapter was created with the correct engine
             mock_adapter_class.assert_called_once_with(mock_engine)
@@ -122,30 +123,26 @@ class TestPermissionService:
         # Mock users
         mock_get_users.return_value = mock_users
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service._maybe_initialize_policies()
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service._maybe_initialize_policies()
 
-            # Verify policies were added for user role blocklist
-            mock_enforcer.add_policy.assert_any_call('user',
-                                                     '/workspaces/config',
-                                                     'POST')
+        # Verify policies were added for user role blocklist
+        mock_enforcer.add_policy.assert_any_call('user', '/workspaces/config',
+                                                 'POST')
 
-            # Verify workspace policies were added
-            mock_enforcer.add_policy.assert_any_call('*', 'default', '*')
-            mock_enforcer.add_policy.assert_any_call('user1', 'private-ws', '*')
-            mock_enforcer.add_policy.assert_any_call('user2', 'private-ws', '*')
+        # Verify workspace policies were added
+        mock_enforcer.add_policy.assert_any_call('*', 'default', '*')
+        mock_enforcer.add_policy.assert_any_call('user1', 'private-ws', '*')
+        mock_enforcer.add_policy.assert_any_call('user2', 'private-ws', '*')
 
-            # Verify users were assigned default roles
-            for user in mock_users:
-                mock_enforcer.add_grouping_policy.assert_any_call(
-                    user.id, rbac.get_default_role())
+        # Verify users were assigned default roles
+        for user in mock_users:
+            mock_enforcer.add_grouping_policy.assert_any_call(
+                user.id, rbac.get_default_role())
 
-            # Verify policy was saved
-            mock_enforcer.save_policy.assert_called()
+        # Verify policy was saved
+        mock_enforcer.save_policy.assert_called()
 
     @mock.patch('sky.users.permission._policy_lock')
     @mock.patch('sky.users.rbac.get_workspace_policy_permissions')
@@ -198,17 +195,14 @@ class TestPermissionService:
         # Mock users
         mock_get_users.return_value = mock_users
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service._maybe_initialize_policies()
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service._maybe_initialize_policies()
 
-            # Verify no new policies were added (since they already exist)
-            mock_enforcer.remove_filtered_policy.assert_not_called()
-            # save_policy should not be called if no updates were made
-            # (users already have roles, policies already exist)
+        # Verify no new policies were added (since they already exist)
+        mock_enforcer.remove_filtered_policy.assert_not_called()
+        # save_policy should not be called if no updates were made
+        # (users already have roles, policies already exist)
 
     def test_add_user_if_not_exists_new_user(self):
         """Test adding a new user that doesn't exist."""
@@ -216,17 +210,14 @@ class TestPermissionService:
         mock_enforcer.get_roles_for_user.return_value = []  # No existing roles
         mock_enforcer.add_grouping_policy.return_value = True
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            result = service._add_user_if_not_exists_no_lock('new_user')
+        result = service._add_user_if_not_exists_no_lock('new_user')
 
-            assert result is True
-            mock_enforcer.add_grouping_policy.assert_called_once_with(
-                'new_user', rbac.get_default_role())
+        assert result is True
+        mock_enforcer.add_grouping_policy.assert_called_once_with(
+            'new_user', rbac.get_default_role())
 
     def test_add_user_if_not_exists_existing_user(self):
         """Test adding a user that already exists."""
@@ -235,16 +226,13 @@ class TestPermissionService:
             'user'
         ]  # User already has roles
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            result = service._add_user_if_not_exists_no_lock('existing_user')
+        result = service._add_user_if_not_exists_no_lock('existing_user')
 
-            assert result is False
-            mock_enforcer.add_grouping_policy.assert_not_called()
+        assert result is False
+        mock_enforcer.add_grouping_policy.assert_not_called()
 
     @mock.patch('sky.users.permission._policy_lock')
     def test_update_role_new_role(self, mock_policy_lock):
@@ -257,21 +245,18 @@ class TestPermissionService:
         mock_enforcer.remove_grouping_policy.return_value = True
         mock_enforcer.add_grouping_policy.return_value = True
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service._load_policy_no_lock = mock.Mock()
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service._load_policy_no_lock = mock.Mock()
 
-            service.update_role('user1', 'admin')
+        service.update_role('user1', 'admin')
 
-            # Verify old role was removed and new role was added
-            mock_enforcer.remove_grouping_policy.assert_called_once_with(
-                'user1', 'user')
-            mock_enforcer.add_grouping_policy.assert_called_once_with(
-                'user1', 'admin')
-            mock_enforcer.save_policy.assert_called_once()
+        # Verify old role was removed and new role was added
+        mock_enforcer.remove_grouping_policy.assert_called_once_with(
+            'user1', 'user')
+        mock_enforcer.add_grouping_policy.assert_called_once_with(
+            'user1', 'admin')
+        mock_enforcer.save_policy.assert_called_once()
 
     @mock.patch('sky.users.permission._policy_lock')
     def test_update_role_same_role(self, mock_policy_lock):
@@ -283,54 +268,44 @@ class TestPermissionService:
         mock_enforcer.get_roles_for_user.return_value = ['admin'
                                                         ]  # Current role
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service._load_policy_no_lock = mock.Mock()
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service._load_policy_no_lock = mock.Mock()
 
-            service.update_role('user1', 'admin')
+        service.update_role('user1', 'admin')
 
-            # Verify no changes were made
-            mock_enforcer.remove_grouping_policy.assert_not_called()
-            mock_enforcer.add_grouping_policy.assert_not_called()
-            mock_enforcer.save_policy.assert_not_called()
+        # Verify no changes were made
+        mock_enforcer.remove_grouping_policy.assert_not_called()
+        mock_enforcer.add_grouping_policy.assert_not_called()
+        mock_enforcer.save_policy.assert_not_called()
 
     def test_get_user_roles(self):
         """Test getting user roles."""
         mock_enforcer = mock.Mock()
         mock_enforcer.get_roles_for_user.return_value = ['admin', 'user']
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service._load_policy_no_lock = mock.Mock()
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service._load_policy_no_lock = mock.Mock()
 
-            roles = service.get_user_roles('user1')
+        roles = service.get_user_roles('user1')
 
-            assert roles == ['admin', 'user']
-            mock_enforcer.get_roles_for_user.assert_called_once_with('user1')
+        assert roles == ['admin', 'user']
+        mock_enforcer.get_roles_for_user.assert_called_once_with('user1')
 
     def test_check_endpoint_permission(self):
         """Test checking endpoint permissions."""
         mock_enforcer = mock.Mock()
         mock_enforcer.enforce.return_value = True
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            result = service.check_endpoint_permission('user1', '/api/test',
-                                                       'GET')
+        result = service.check_endpoint_permission('user1', '/api/test', 'GET')
 
-            assert result is True
-            mock_enforcer.enforce.assert_called_once_with(
-                'user1', '/api/test', 'GET')
+        assert result is True
+        mock_enforcer.enforce.assert_called_once_with('user1', '/api/test',
+                                                      'GET')
 
     def test_check_workspace_permission_non_server(self):
         """Test workspace permission check when not on API server."""
@@ -338,16 +313,13 @@ class TestPermissionService:
         if constants.ENV_VAR_IS_SKYPILOT_SERVER in os.environ:
             del os.environ[constants.ENV_VAR_IS_SKYPILOT_SERVER]
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
+        service = permission.PermissionService()
+        service._lazy_initialize()
 
-            result = service.check_workspace_permission('user1',
-                                                        'test-workspace')
+        result = service.check_workspace_permission('user1', 'test-workspace')
 
-            # Should always return True when not on API server
-            assert result is True
+        # Should always return True when not on API server
+        assert result is True
 
     def test_check_workspace_permission_admin_user(self):
         """Test workspace permission check for admin user."""
@@ -355,18 +327,15 @@ class TestPermissionService:
 
         mock_enforcer = mock.Mock()
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service.get_user_roles = mock.Mock(return_value=['admin'])
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service.get_user_roles = mock.Mock(return_value=['admin'])
 
-            result = service.check_workspace_permission('admin_user',
-                                                        'test-workspace')
+        result = service.check_workspace_permission('admin_user',
+                                                    'test-workspace')
 
-            # Admin should always have access
-            assert result is True
+        # Admin should always have access
+        assert result is True
 
     def test_check_workspace_permission_regular_user(self):
         """Test workspace permission check for regular user."""
@@ -375,19 +344,15 @@ class TestPermissionService:
         mock_enforcer = mock.Mock()
         mock_enforcer.enforce.return_value = True
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service.get_user_roles = mock.Mock(return_value=['user'])
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service.get_user_roles = mock.Mock(return_value=['user'])
 
-            result = service.check_workspace_permission('user1',
-                                                        'test-workspace')
+        result = service.check_workspace_permission('user1', 'test-workspace')
 
-            assert result is True
-            mock_enforcer.enforce.assert_called_once_with(
-                'user1', 'test-workspace', '*')
+        assert result is True
+        mock_enforcer.enforce.assert_called_once_with('user1', 'test-workspace',
+                                                      '*')
 
     @mock.patch('sky.users.permission._policy_lock')
     def test_add_workspace_policy(self, mock_policy_lock):
@@ -398,20 +363,15 @@ class TestPermissionService:
         mock_enforcer = mock.Mock()
         mock_enforcer.add_policy.return_value = True
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            service.add_workspace_policy('test-workspace', ['user1', 'user2'])
+        service.add_workspace_policy('test-workspace', ['user1', 'user2'])
 
-            # Verify policies were added for each user
-            mock_enforcer.add_policy.assert_any_call('user1', 'test-workspace',
-                                                     '*')
-            mock_enforcer.add_policy.assert_any_call('user2', 'test-workspace',
-                                                     '*')
-            mock_enforcer.save_policy.assert_called_once()
+        # Verify policies were added for each user
+        mock_enforcer.add_policy.assert_any_call('user1', 'test-workspace', '*')
+        mock_enforcer.add_policy.assert_any_call('user2', 'test-workspace', '*')
+        mock_enforcer.save_policy.assert_called_once()
 
     @mock.patch('sky.users.permission.filelock.FileLock')
     def test_policy_lock_context_manager(self, mock_filelock):
@@ -472,9 +432,11 @@ class TestPermissionServiceMultiProcess:
                                '_maybe_initialize_policies'):
             # Create first instance
             service1 = permission.PermissionService()
+            service1._lazy_initialize()
 
             # Create second instance
             service2 = permission.PermissionService()
+            service2._lazy_initialize()
 
             # Both should share the same enforcer
             assert service1.enforcer is service2.enforcer
@@ -541,13 +503,10 @@ class TestPermissionServiceMultiProcess:
 
         def create_service():
             try:
-                with mock.patch.object(permission.PermissionService,
-                                       '__init__',
-                                       return_value=None):
-                    service = permission.PermissionService()
-                    service.enforcer = mock_enforcer
-                    service._maybe_initialize_policies()
-                    services.append(service)
+                service = permission.PermissionService()
+                service.enforcer = mock_enforcer
+                service._maybe_initialize_policies()
+                services.append(service)
             except Exception as e:
                 errors.append(e)
 
@@ -636,26 +595,23 @@ class TestPermissionServiceMultiProcess:
         mock_get_workspace_perms.return_value = {}
         mock_get_users.return_value = mock_users
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            # Initialize policies multiple times
-            service._maybe_initialize_policies()
-            service._maybe_initialize_policies()
-            service._maybe_initialize_policies()
+        # Initialize policies multiple times
+        service._maybe_initialize_policies()
+        service._maybe_initialize_policies()
+        service._maybe_initialize_policies()
 
-            # Each user should only be added once (3 users total)
-            assert len(grouping_policy_calls) == len(mock_users)
+        # Each user should only be added once (3 users total)
+        assert len(grouping_policy_calls) == len(mock_users)
 
-            # Verify each user was added exactly once
-            expected_calls = {
-                (user.id, rbac.get_default_role()) for user in mock_users
-            }
-            actual_calls = set(grouping_policy_calls)
-            assert actual_calls == expected_calls
+        # Verify each user was added exactly once
+        expected_calls = {
+            (user.id, rbac.get_default_role()) for user in mock_users
+        }
+        actual_calls = set(grouping_policy_calls)
+        assert actual_calls == expected_calls
 
     @mock.patch('sky.users.permission._policy_lock')
     def test_concurrent_policy_updates_use_lock(self, mock_policy_lock):
@@ -680,20 +636,17 @@ class TestPermissionServiceMultiProcess:
         mock_enforcer.add_policy.return_value = True
         mock_enforcer.save_policy.return_value = True
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            # Test a single workspace policy update to verify lock is used
-            service.add_workspace_policy('test-workspace', ['user1'])
+        # Test a single workspace policy update to verify lock is used
+        service.add_workspace_policy('test-workspace', ['user1'])
 
-            # Verify that lock was acquired and released
-            assert 'enter' in lock_calls
-            assert 'exit' in lock_calls
-            assert len([call for call in lock_calls if call == 'enter']) == 1
-            assert len([call for call in lock_calls if call == 'exit']) == 1
+        # Verify that lock was acquired and released
+        assert 'enter' in lock_calls
+        assert 'exit' in lock_calls
+        assert len([call for call in lock_calls if call == 'enter']) == 1
+        assert len([call for call in lock_calls if call == 'exit']) == 1
 
     @mock.patch('sky.users.permission._policy_lock')
     @mock.patch('sky.users.rbac.get_workspace_policy_permissions')
@@ -737,18 +690,15 @@ class TestPermissionServiceMultiProcess:
 
         mock_get_users.return_value = mock_users
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            # Initialize policies when they already exist
-            service._maybe_initialize_policies()
+        # Initialize policies when they already exist
+        service._maybe_initialize_policies()
 
-            # Since policies already exist and users already have roles,
-            # no new policies or grouping policies should be added
-            mock_enforcer.remove_filtered_policy.assert_not_called()
+        # Since policies already exist and users already have roles,
+        # no new policies or grouping policies should be added
+        mock_enforcer.remove_filtered_policy.assert_not_called()
 
-            # The save_policy might not be called if no changes were made
-            # We mainly want to ensure no duplicate additions occurred
+        # The save_policy might not be called if no changes were made
+        # We mainly want to ensure no duplicate additions occurred