Fix #9 tests

conwetlab · Oct 18, 2016 · 4601615 · 4601615
1 parent f6959c4
commit 4601615
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/ckanext/oauth2/tests/test_controller.py b/ckanext/oauth2/tests/test_controller.py
@@ -57,13 +57,17 @@ def setUp(self):
         self._toolkit = controller.toolkit
         controller.toolkit = MagicMock()
 
+        self.__session = controller.session
+        controller.session = MagicMock()
+
         self.controller = controller.OAuth2Controller()
 
     def tearDown(self):
         # Unmock the function
         controller.helpers = self._helpers
         controller.oauth2 = self._oauth2
         controller.toolkit = self._toolkit
+        controller.session = self.__session
 
     def generate_state(self, url):
         return b64encode(bytes(json.dumps({CAME_FROM_FIELD: url})))
@@ -118,6 +122,7 @@ def test_controller_errors(self, came_from=None, exception=Exception(EXCEPTION_M
         self.controller.callback()
 
         # Check the state and the location
+        controller.session.save.assert_called_once_with()
         self.assertEquals(RETURNED_STATUS, controller.toolkit.response.status_int)
         self.assertEquals(came_from, controller.toolkit.response.location)
         controller.helpers.flash_error.assert_called_once_with(expected_flash)
diff --git a/ckanext/oauth2/tests/test_plugin.py b/ckanext/oauth2/tests/test_plugin.py
@@ -37,8 +37,6 @@ def setUp(self):
         self._toolkit = plugin.toolkit
         plugin.toolkit = MagicMock()
 
-        self._session = plugin.session
-        plugin.session = MagicMock()
 
         self._oauth2 = plugin.oauth2
         plugin.oauth2 = MagicMock()
@@ -50,7 +48,6 @@ def tearDown(self):
         # Unmock functions
         plugin.config = self._config
         plugin.toolkit = self._toolkit
-        plugin.session = self._session
 
     def _set_identity(self, identity):
         plugin.toolkit.request.environ = {}
@@ -166,7 +163,6 @@ def authenticate_side_effect(identity):
             self.assertEquals(0, plugin.oauth2.OAuth2Helper.return_value.identify.call_count)
 
         self.assertEquals(expected_user, plugin.toolkit.c.user)
-        plugin.session.save.assert_called_once()
 
         if expected_user is None:
             self.assertIsNone(plugin.toolkit.c.usertoken)