| # Copyright 2015 Google Inc. All rights reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| """Unit tests for oauth2client.multistore_file.""" |
| |
| import contextlib |
| import datetime |
| import json |
| import multiprocessing |
| import os |
| import tempfile |
| |
| import fasteners |
| import mock |
| from six import StringIO |
| import unittest2 |
| |
| from oauth2client import client |
| from oauth2client.contrib import multiprocess_file_storage |
| |
| from ..http_mock import HttpMockSequence |
| |
| |
| @contextlib.contextmanager |
| def scoped_child_process(target, **kwargs): |
| die_event = multiprocessing.Event() |
| ready_event = multiprocessing.Event() |
| process = multiprocessing.Process( |
| target=target, args=(die_event, ready_event), kwargs=kwargs) |
| process.start() |
| try: |
| ready_event.wait() |
| yield |
| finally: |
| die_event.set() |
| process.join(5) |
| |
| |
| def _create_test_credentials(expiration=None): |
| access_token = 'foo' |
| client_secret = 'cOuDdkfjxxnv+' |
| refresh_token = '1/0/a.df219fjls0' |
| token_expiry = expiration or ( |
| datetime.datetime.utcnow() + datetime.timedelta(seconds=3600)) |
| token_uri = 'https://www.google.com/accounts/o8/oauth2/token' |
| user_agent = 'refresh_checker/1.0' |
| |
| credentials = client.OAuth2Credentials( |
| access_token, 'test-client-id', client_secret, |
| refresh_token, token_expiry, token_uri, |
| user_agent) |
| return credentials |
| |
| |
| def _generate_token_response_http(new_token='new_token'): |
| token_response = json.dumps({ |
| 'access_token': new_token, |
| 'expires_in': '3600', |
| }) |
| http = HttpMockSequence([ |
| ({'status': '200'}, token_response), |
| ]) |
| |
| return http |
| |
| |
| class MultiprocessStorageBehaviorTests(unittest2.TestCase): |
| |
| def setUp(self): |
| filehandle, self.filename = tempfile.mkstemp( |
| 'oauth2client_test.data') |
| os.close(filehandle) |
| |
| def tearDown(self): |
| try: |
| os.unlink(self.filename) |
| os.unlink('{0}.lock'.format(self.filename)) |
| except OSError: # pragma: NO COVER |
| pass |
| |
| def test_basic_operations(self): |
| credentials = _create_test_credentials() |
| |
| store = multiprocess_file_storage.MultiprocessFileStorage( |
| self.filename, 'basic') |
| |
| # Save credentials |
| store.put(credentials) |
| credentials = store.get() |
| |
| self.assertIsNotNone(credentials) |
| self.assertEqual('foo', credentials.access_token) |
| |
| # Reset internal cache, ensure credentials were saved. |
| store._backend._credentials = {} |
| credentials = store.get() |
| |
| self.assertIsNotNone(credentials) |
| self.assertEqual('foo', credentials.access_token) |
| |
| # Delete credentials |
| store.delete() |
| credentials = store.get() |
| |
| self.assertIsNone(credentials) |
| |
| def test_single_process_refresh(self): |
| store = multiprocess_file_storage.MultiprocessFileStorage( |
| self.filename, 'single-process') |
| credentials = _create_test_credentials() |
| credentials.set_store(store) |
| |
| http = _generate_token_response_http() |
| credentials.refresh(http) |
| self.assertEqual(credentials.access_token, 'new_token') |
| |
| retrieved = store.get() |
| self.assertEqual(retrieved.access_token, 'new_token') |
| |
| def test_multi_process_refresh(self): |
| # This will test that two processes attempting to refresh credentials |
| # will only refresh once. |
| store = multiprocess_file_storage.MultiprocessFileStorage( |
| self.filename, 'multi-process') |
| credentials = _create_test_credentials() |
| credentials.set_store(store) |
| store.put(credentials) |
| |
| def child_process_func( |
| die_event, ready_event, check_event): # pragma: NO COVER |
| store = multiprocess_file_storage.MultiprocessFileStorage( |
| self.filename, 'multi-process') |
| |
| credentials = store.get() |
| self.assertIsNotNone(credentials) |
| |
| # Make sure this thread gets to refresh first. |
| original_acquire_lock = store.acquire_lock |
| |
| def replacement_acquire_lock(*args, **kwargs): |
| result = original_acquire_lock(*args, **kwargs) |
| ready_event.set() |
| check_event.wait() |
| return result |
| |
| credentials.store.acquire_lock = replacement_acquire_lock |
| |
| http = _generate_token_response_http('b') |
| credentials.refresh(http) |
| |
| self.assertEqual(credentials.access_token, 'b') |
| |
| check_event = multiprocessing.Event() |
| with scoped_child_process(child_process_func, check_event=check_event): |
| # The lock should be currently held by the child process. |
| self.assertFalse( |
| store._backend._process_lock.acquire(blocking=False)) |
| check_event.set() |
| |
| # The child process will refresh first, so we should end up |
| # with 'b' as the token. |
| http = mock.Mock() |
| credentials.refresh(http=http) |
| self.assertEqual(credentials.access_token, 'b') |
| self.assertFalse(http.request.called) |
| |
| retrieved = store.get() |
| self.assertEqual(retrieved.access_token, 'b') |
| |
| def test_read_only_file_fail_lock(self): |
| credentials = _create_test_credentials() |
| |
| # Grab the lock in another process, preventing this process from |
| # acquiring the lock. |
| def child_process(die_event, ready_event): # pragma: NO COVER |
| lock = fasteners.InterProcessLock( |
| '{0}.lock'.format(self.filename)) |
| with lock: |
| ready_event.set() |
| die_event.wait() |
| |
| with scoped_child_process(child_process): |
| store = multiprocess_file_storage.MultiprocessFileStorage( |
| self.filename, 'fail-lock') |
| store.put(credentials) |
| self.assertTrue(store._backend._read_only) |
| |
| # These credentials should still be in the store's memory-only cache. |
| self.assertIsNotNone(store.get()) |
| |
| |
| class MultiprocessStorageUnitTests(unittest2.TestCase): |
| |
| def setUp(self): |
| filehandle, self.filename = tempfile.mkstemp( |
| 'oauth2client_test.data') |
| os.close(filehandle) |
| |
| def tearDown(self): |
| try: |
| os.unlink(self.filename) |
| os.unlink('{0}.lock'.format(self.filename)) |
| except OSError: # pragma: NO COVER |
| pass |
| |
| def test__create_file_if_needed(self): |
| self.assertFalse( |
| multiprocess_file_storage._create_file_if_needed(self.filename)) |
| os.unlink(self.filename) |
| self.assertTrue( |
| multiprocess_file_storage._create_file_if_needed(self.filename)) |
| self.assertTrue( |
| os.path.exists(self.filename)) |
| |
| def test__get_backend(self): |
| backend_one = multiprocess_file_storage._get_backend('file_a') |
| backend_two = multiprocess_file_storage._get_backend('file_a') |
| backend_three = multiprocess_file_storage._get_backend('file_b') |
| |
| self.assertIs(backend_one, backend_two) |
| self.assertIsNot(backend_one, backend_three) |
| |
| def test__read_write_credentials_file(self): |
| credentials = _create_test_credentials() |
| contents = StringIO() |
| |
| multiprocess_file_storage._write_credentials_file( |
| contents, {'key': credentials}) |
| |
| contents.seek(0) |
| data = json.load(contents) |
| self.assertEqual(data['file_version'], 2) |
| self.assertTrue(data['credentials']['key']) |
| |
| # Read it back. |
| contents.seek(0) |
| results = multiprocess_file_storage._load_credentials_file(contents) |
| self.assertEqual( |
| results['key'].access_token, credentials.access_token) |
| |
| # Add an invalid credential and try reading it back. It should ignore |
| # the invalid one but still load the valid one. |
| data['credentials']['invalid'] = '123' |
| results = multiprocess_file_storage._load_credentials_file( |
| StringIO(json.dumps(data))) |
| self.assertNotIn('invalid', results) |
| self.assertEqual( |
| results['key'].access_token, credentials.access_token) |
| |
| def test__load_credentials_file_invalid_json(self): |
| contents = StringIO('{[') |
| self.assertEqual( |
| multiprocess_file_storage._load_credentials_file(contents), {}) |
| |
| def test__load_credentials_file_no_file_version(self): |
| contents = StringIO('{}') |
| self.assertEqual( |
| multiprocess_file_storage._load_credentials_file(contents), {}) |
| |
| def test__load_credentials_file_bad_file_version(self): |
| contents = StringIO(json.dumps({'file_version': 1})) |
| self.assertEqual( |
| multiprocess_file_storage._load_credentials_file(contents), {}) |
| |
| def test__load_credentials_no_open_file(self): |
| backend = multiprocess_file_storage._get_backend(self.filename) |
| backend._credentials = mock.Mock() |
| backend._credentials.update.side_effect = AssertionError() |
| backend._load_credentials() |
| |
| def test_acquire_lock_nonexistent_file(self): |
| backend = multiprocess_file_storage._get_backend(self.filename) |
| os.unlink(self.filename) |
| backend._process_lock = mock.Mock() |
| backend._process_lock.acquire.return_value = False |
| backend.acquire_lock() |
| self.assertIsNone(backend._file) |
| |
| def test_release_lock_with_no_file(self): |
| backend = multiprocess_file_storage._get_backend(self.filename) |
| backend._file = None |
| backend._read_only = True |
| backend._thread_lock.acquire() |
| backend.release_lock() |
| |
| def test__refresh_predicate(self): |
| backend = multiprocess_file_storage._get_backend(self.filename) |
| |
| credentials = _create_test_credentials() |
| self.assertFalse(backend._refresh_predicate(credentials)) |
| |
| credentials.invalid = True |
| self.assertTrue(backend._refresh_predicate(credentials)) |
| |
| credentials = _create_test_credentials( |
| expiration=( |
| datetime.datetime.utcnow() - datetime.timedelta(seconds=3600))) |
| self.assertTrue(backend._refresh_predicate(credentials)) |
| |
| |
| if __name__ == '__main__': # pragma: NO COVER |
| unittest2.main() |
