# 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()