265 lines
9.2 KiB
Python
265 lines
9.2 KiB
Python
|
# Copyright 2016 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.
|
||
|
|
||
|
|
||
|
import unittest2
|
||
|
|
||
|
|
||
|
class TestConnectionPool(unittest2.TestCase):
|
||
|
|
||
|
def _getTargetClass(self):
|
||
|
from gcloud.bigtable.happybase.pool import ConnectionPool
|
||
|
return ConnectionPool
|
||
|
|
||
|
def _makeOne(self, *args, **kwargs):
|
||
|
return self._getTargetClass()(*args, **kwargs)
|
||
|
|
||
|
def test_constructor_defaults(self):
|
||
|
import six
|
||
|
import threading
|
||
|
from gcloud.bigtable.happybase.connection import Connection
|
||
|
|
||
|
size = 11
|
||
|
instance_copy = _Instance()
|
||
|
all_copies = [instance_copy] * size
|
||
|
instance = _Instance(all_copies) # Avoid implicit environ check.
|
||
|
pool = self._makeOne(size, instance=instance)
|
||
|
|
||
|
self.assertTrue(isinstance(pool._lock, type(threading.Lock())))
|
||
|
self.assertTrue(isinstance(pool._thread_connections, threading.local))
|
||
|
self.assertEqual(pool._thread_connections.__dict__, {})
|
||
|
|
||
|
queue = pool._queue
|
||
|
self.assertTrue(isinstance(queue, six.moves.queue.LifoQueue))
|
||
|
self.assertTrue(queue.full())
|
||
|
self.assertEqual(queue.maxsize, size)
|
||
|
for connection in queue.queue:
|
||
|
self.assertTrue(isinstance(connection, Connection))
|
||
|
self.assertTrue(connection._instance is instance_copy)
|
||
|
|
||
|
def test_constructor_passes_kwargs(self):
|
||
|
table_prefix = 'foo'
|
||
|
table_prefix_separator = '<>'
|
||
|
instance = _Instance() # Avoid implicit environ check.
|
||
|
|
||
|
size = 1
|
||
|
pool = self._makeOne(size, table_prefix=table_prefix,
|
||
|
table_prefix_separator=table_prefix_separator,
|
||
|
instance=instance)
|
||
|
|
||
|
for connection in pool._queue.queue:
|
||
|
self.assertEqual(connection.table_prefix, table_prefix)
|
||
|
self.assertEqual(connection.table_prefix_separator,
|
||
|
table_prefix_separator)
|
||
|
|
||
|
def test_constructor_ignores_autoconnect(self):
|
||
|
from gcloud._testing import _Monkey
|
||
|
from gcloud.bigtable.happybase.connection import Connection
|
||
|
from gcloud.bigtable.happybase import pool as MUT
|
||
|
|
||
|
class ConnectionWithOpen(Connection):
|
||
|
|
||
|
_open_called = False
|
||
|
|
||
|
def open(self):
|
||
|
self._open_called = True
|
||
|
|
||
|
# First make sure the custom Connection class does as expected.
|
||
|
instance_copy1 = _Instance()
|
||
|
instance_copy2 = _Instance()
|
||
|
instance_copy3 = _Instance()
|
||
|
instance = _Instance([instance_copy1, instance_copy2, instance_copy3])
|
||
|
connection = ConnectionWithOpen(autoconnect=False, instance=instance)
|
||
|
self.assertFalse(connection._open_called)
|
||
|
self.assertTrue(connection._instance is instance_copy1)
|
||
|
connection = ConnectionWithOpen(autoconnect=True, instance=instance)
|
||
|
self.assertTrue(connection._open_called)
|
||
|
self.assertTrue(connection._instance is instance_copy2)
|
||
|
|
||
|
# Then make sure autoconnect=True is ignored in a pool.
|
||
|
size = 1
|
||
|
with _Monkey(MUT, Connection=ConnectionWithOpen):
|
||
|
pool = self._makeOne(size, autoconnect=True, instance=instance)
|
||
|
|
||
|
for connection in pool._queue.queue:
|
||
|
self.assertTrue(isinstance(connection, ConnectionWithOpen))
|
||
|
self.assertTrue(connection._instance is instance_copy3)
|
||
|
self.assertFalse(connection._open_called)
|
||
|
|
||
|
def test_constructor_infers_instance(self):
|
||
|
from gcloud._testing import _Monkey
|
||
|
from gcloud.bigtable.happybase.connection import Connection
|
||
|
from gcloud.bigtable.happybase import pool as MUT
|
||
|
|
||
|
size = 1
|
||
|
instance_copy = _Instance()
|
||
|
all_copies = [instance_copy] * size
|
||
|
instance = _Instance(all_copies)
|
||
|
get_instance_calls = []
|
||
|
|
||
|
def mock_get_instance(timeout=None):
|
||
|
get_instance_calls.append(timeout)
|
||
|
return instance
|
||
|
|
||
|
with _Monkey(MUT, _get_instance=mock_get_instance):
|
||
|
pool = self._makeOne(size)
|
||
|
|
||
|
for connection in pool._queue.queue:
|
||
|
self.assertTrue(isinstance(connection, Connection))
|
||
|
# We know that the Connection() constructor will
|
||
|
# call instance.copy().
|
||
|
self.assertTrue(connection._instance is instance_copy)
|
||
|
|
||
|
self.assertEqual(get_instance_calls, [None])
|
||
|
|
||
|
def test_constructor_non_integer_size(self):
|
||
|
size = None
|
||
|
with self.assertRaises(TypeError):
|
||
|
self._makeOne(size)
|
||
|
|
||
|
def test_constructor_non_positive_size(self):
|
||
|
size = -10
|
||
|
with self.assertRaises(ValueError):
|
||
|
self._makeOne(size)
|
||
|
size = 0
|
||
|
with self.assertRaises(ValueError):
|
||
|
self._makeOne(size)
|
||
|
|
||
|
def _makeOneWithMockQueue(self, queue_return):
|
||
|
from gcloud._testing import _Monkey
|
||
|
from gcloud.bigtable.happybase import pool as MUT
|
||
|
|
||
|
# We are going to use a fake queue, so we don't want any connections
|
||
|
# or instances to be created in the constructor.
|
||
|
size = -1
|
||
|
instance = object()
|
||
|
with _Monkey(MUT, _MIN_POOL_SIZE=size):
|
||
|
pool = self._makeOne(size, instance=instance)
|
||
|
|
||
|
pool._queue = _Queue(queue_return)
|
||
|
return pool
|
||
|
|
||
|
def test__acquire_connection(self):
|
||
|
queue_return = object()
|
||
|
pool = self._makeOneWithMockQueue(queue_return)
|
||
|
|
||
|
timeout = 432
|
||
|
connection = pool._acquire_connection(timeout=timeout)
|
||
|
self.assertTrue(connection is queue_return)
|
||
|
self.assertEqual(pool._queue._get_calls, [(True, timeout)])
|
||
|
self.assertEqual(pool._queue._put_calls, [])
|
||
|
|
||
|
def test__acquire_connection_failure(self):
|
||
|
from gcloud.bigtable.happybase.pool import NoConnectionsAvailable
|
||
|
|
||
|
pool = self._makeOneWithMockQueue(None)
|
||
|
timeout = 1027
|
||
|
with self.assertRaises(NoConnectionsAvailable):
|
||
|
pool._acquire_connection(timeout=timeout)
|
||
|
self.assertEqual(pool._queue._get_calls, [(True, timeout)])
|
||
|
self.assertEqual(pool._queue._put_calls, [])
|
||
|
|
||
|
def test_connection_is_context_manager(self):
|
||
|
import contextlib
|
||
|
import six
|
||
|
|
||
|
queue_return = _Connection()
|
||
|
pool = self._makeOneWithMockQueue(queue_return)
|
||
|
cnxn_context = pool.connection()
|
||
|
if six.PY3: # pragma: NO COVER Python 3
|
||
|
self.assertTrue(isinstance(cnxn_context,
|
||
|
contextlib._GeneratorContextManager))
|
||
|
else:
|
||
|
self.assertTrue(isinstance(cnxn_context,
|
||
|
contextlib.GeneratorContextManager))
|
||
|
|
||
|
def test_connection_no_current_cnxn(self):
|
||
|
queue_return = _Connection()
|
||
|
pool = self._makeOneWithMockQueue(queue_return)
|
||
|
timeout = 55
|
||
|
|
||
|
self.assertFalse(hasattr(pool._thread_connections, 'current'))
|
||
|
with pool.connection(timeout=timeout) as connection:
|
||
|
self.assertEqual(pool._thread_connections.current, queue_return)
|
||
|
self.assertTrue(connection is queue_return)
|
||
|
self.assertFalse(hasattr(pool._thread_connections, 'current'))
|
||
|
|
||
|
self.assertEqual(pool._queue._get_calls, [(True, timeout)])
|
||
|
self.assertEqual(pool._queue._put_calls,
|
||
|
[(queue_return, None, None)])
|
||
|
|
||
|
def test_connection_with_current_cnxn(self):
|
||
|
current_cnxn = _Connection()
|
||
|
queue_return = _Connection()
|
||
|
pool = self._makeOneWithMockQueue(queue_return)
|
||
|
pool._thread_connections.current = current_cnxn
|
||
|
timeout = 8001
|
||
|
|
||
|
with pool.connection(timeout=timeout) as connection:
|
||
|
self.assertTrue(connection is current_cnxn)
|
||
|
|
||
|
self.assertEqual(pool._queue._get_calls, [])
|
||
|
self.assertEqual(pool._queue._put_calls, [])
|
||
|
self.assertEqual(pool._thread_connections.current, current_cnxn)
|
||
|
|
||
|
|
||
|
class _Client(object):
|
||
|
|
||
|
def __init__(self):
|
||
|
self.stop_calls = 0
|
||
|
|
||
|
def stop(self):
|
||
|
self.stop_calls += 1
|
||
|
|
||
|
|
||
|
class _Connection(object):
|
||
|
|
||
|
def open(self):
|
||
|
pass
|
||
|
|
||
|
|
||
|
class _Instance(object):
|
||
|
|
||
|
def __init__(self, copies=()):
|
||
|
self.copies = list(copies)
|
||
|
# Included to support Connection.__del__
|
||
|
self._client = _Client()
|
||
|
|
||
|
def copy(self):
|
||
|
if self.copies:
|
||
|
result = self.copies[0]
|
||
|
self.copies[:] = self.copies[1:]
|
||
|
return result
|
||
|
else:
|
||
|
return self
|
||
|
|
||
|
|
||
|
class _Queue(object):
|
||
|
|
||
|
def __init__(self, result=None):
|
||
|
self.result = result
|
||
|
self._get_calls = []
|
||
|
self._put_calls = []
|
||
|
|
||
|
def get(self, block=None, timeout=None):
|
||
|
self._get_calls.append((block, timeout))
|
||
|
if self.result is None:
|
||
|
import six
|
||
|
raise six.moves.queue.Empty
|
||
|
else:
|
||
|
return self.result
|
||
|
|
||
|
def put(self, item, block=None, timeout=None):
|
||
|
self._put_calls.append((item, block, timeout))
|