diff --git a/cloudbaseinit/tests/plugins/windows/test_winrmlistener.py b/cloudbaseinit/tests/plugins/windows/test_winrmlistener.py
index 22092f01e..f14748617 100644
--- a/cloudbaseinit/tests/plugins/windows/test_winrmlistener.py
+++ b/cloudbaseinit/tests/plugins/windows/test_winrmlistener.py
@@ -102,9 +102,12 @@ def _test_check_uac_remote_restrictions(self, mock_SecurityUtils,
         mock_SecurityUtils.return_value = mock_security_utils
         mock_osutils = mock.Mock()
         mock_osutils.check_os_version.side_effect = [True, False]
+        mock_security_utils.get_uac_remote_restrictions.return_value = \
+            disable_uac_remote_restrictions
+        expected_set_token_calls = []
         if disable_uac_remote_restrictions:
-            mock_security_utils.get_uac_remote_restrictions.return_value = \
-                disable_uac_remote_restrictions
+            expected_set_token_calls = [mock.call(enable=False),
+                                        mock.call(enable=True)]
 
         with self._winrmlistener._check_uac_remote_restrictions(mock_osutils):
             mock_SecurityUtils.assert_called_once_with()
@@ -112,13 +115,9 @@ def _test_check_uac_remote_restrictions(self, mock_SecurityUtils,
                 [mock.call(6, 0), mock.call(6, 2)])
             (mock_security_utils.get_uac_remote_restrictions.
              assert_called_once_with())
-            if disable_uac_remote_restrictions:
-                expected_set_token_calls = [mock.call(enable=True)]
-            else:
-                expected_set_token_calls = [mock.call(enable=False),
-                                            mock.call(enable=True)]
-            mock_security_utils.set_uac_remote_restrictions.has_calls(
-                expected_set_token_calls)
+
+        mock_security_utils.set_uac_remote_restrictions.assert_has_calls(
+            expected_set_token_calls)
 
     def test_check_uac_remote_restrictions(self):
         self._test_check_uac_remote_restrictions(