Permalink
Browse files

Merge pull request #185 from eventbrite/stubs

Adding stub objects for ColumnFamily, ConnectionPool, and SystemManager
  • Loading branch information...
thobbs committed Jan 31, 2013
2 parents 71ed5b9 + 3820775 commit 957feee52a421242359e4704633e446aeaa95619
View
@@ -18,3 +18,4 @@ Pycassa Modules
pycassa/util
pycassa/logging/pycassa_logger
pycassa/logging/pool_stats_logger
+ pycassa/contrib/stubs
@@ -0,0 +1,34 @@
+:mod:`pycassa.contrib.stubs` -- Pycassa Stubs
+=============================================
+
+.. automodule:: pycassa.contrib.stubs
+
+ .. autoclass:: pycassa.contrib.stubs.ColumnFamilyStub(pool=None, column_family=None, rows=None)
+
+ .. automethod:: get(key[, columns][, include_timestamp])
+
+ .. automethod:: multiget(keys[, columns][, include_timestamp])
+
+ .. automethod:: get_range([columns][, include_timestamp])
+
+ .. automethod:: get_indexed_slices(index_clause[, columns], include_timestamp])
+
+ .. automethod:: insert(key, columns[, timestamp])
+
+ .. automethod:: remove(key[, columns])
+
+ .. automethod:: truncate()
+
+ .. automethod:: batch(self)
+
+ .. autoclass:: pycassa.contrib.stubs.ConnectionPoolStub()
+
+ .. autoclass:: pycassa.contrib.stubs.SystemManagerStub()
+
+ .. automethod:: create_column_family(keyspace, table_name)
+
+ .. automethod:: alter_column(keyspace, table_name, column_name, column_type)
+
+ .. automethod:: create_index(keyspace, table_name, column_name, column_type)
+
+ .. automethod:: describe_schema_versions()
No changes.
View
@@ -0,0 +1,212 @@
+"""A functional set of stubs to be used for unit testing.
+
+Projects that use pycassa and need to run an automated unit test suite on a
+system like Jenkins can use these stubs to emulate interactions with Cassandra
+without spinning up a cluster locally.
+
+"""
+
+import operator
+
+from functools import partial
+
+from pycassa import NotFoundException
+from pycassa.util import OrderedDict
+from pycassa.columnfamily import gm_timestamp
+from pycassa.index import EQ, GT, GTE, LT, LTE
+
+
+__all__ = ['ConnectionPoolStub', 'ColumnFamilyStub', 'SystemManagerStub']
+
+class OrderedDictWithTime(OrderedDict):
+ def __init__(self, *args, **kwargs):
+ self.__timestamp = kwargs.pop('timestamp', None)
+ super(OrderedDictWithTime, self).__init__(*args, **kwargs)
+
+ def __setitem__(self, key, value, timestamp=None):
+ if timestamp is None:
+ timestamp = self.__timestamp or gm_timestamp()
+
+ super(OrderedDictWithTime, self).__setitem__(key, (value, timestamp))
+
+
+operator_dict = {
+ EQ: operator.eq,
+ GT: operator.gt,
+ GTE: operator.ge,
+ LT: operator.lt,
+ LTE: operator.le,
+}
+
+
+class ConnectionPoolStub(object):
+ """Connection pool stub.
+
+ Notes created column families in :attr:`self.column_families`.
+
+ """
+ def __init__(self, *args, **kwargs):
+ self.column_families = {}
+
+ def _register_mock_cf(self, name, cf):
+ if name:
+ self.column_families[name] = cf
+
+ def dispose(self, *args, **kwargs):
+ pass
+
+
+class SystemManagerStub(object):
+ """Functional System Manager stub object.
+
+ Records when column families, columns, and indexes have been created. To
+ see what has been recorded, look at :attr:`self.column_families`.
+
+ """
+
+ def __init__(self, *args, **kwargs):
+ self.column_families = {}
+
+ def create_column_family(self, keyspace, table_name, *args, **kwargs):
+ """Create a column family and record its existence."""
+
+ self.column_families[table_name] = {
+ 'keyspace': keyspace,
+ 'columns': {},
+ 'indexes': {},
+ }
+
+ def alter_column(self, keyspace, table_name, column_name, column_type):
+ """Alter a column, recording its name and type."""
+
+ self.column_families[table_name]['columns'][column_name] = column_type
+
+ def create_index(self, keyspace, table_name, column_name, column_type):
+ """Create an index, recording its name and type."""
+
+ self.column_families[table_name]['indexes'][column_name] = column_type
+
+ def _schema(self):
+ ret = ','.join(self.column_families.keys())
+ for k in self.column_families:
+ for v in ('columns', 'indexes'):
+ ret += ','.join(self.column_families[k][v])
+
+ return hash(ret)
+
+ def describe_schema_versions(self):
+ """Describes the schema based on a hash of the stub system state."""
+
+ return {self._schema(): ['1.1.1.1']}
+
+
+class ColumnFamilyStub(object):
+ """Functional ColumnFamily stub object.
+
+ Acts very similar to a remote column family, supporting a basic version of
+ the API. When instantiated, it registers itself with the supplied (stub)
+ connection pool.
+
+ """
+
+ def __init__(self, pool=None, column_family=None, rows=None, **kwargs):
+ rows = rows or OrderedDict()
+ for r in rows.itervalues():
+ if not isinstance(r, OrderedDictWithTime):
+ r = OrderedDictWithTime(r)
+ self.rows = rows
+
+ if pool is not None:
+ pool._register_mock_cf(column_family, self)
+
+ def __len__(self):
+ return len(self.rows)
+
+ def __contains__(self, obj):
+ return self.rows.__contains__(obj)
+
+ def get(self, key, columns=None, include_timestamp=False, **kwargs):
+ """Get a value from the column family stub."""
+
+ my_columns = self.rows.get(key)
+ if include_timestamp:
+ get_value = lambda x: x
+ else:
+ get_value = lambda x: x[0]
+ if not my_columns:
+ raise NotFoundException()
+
+ return OrderedDict((k, get_value(v)) for (k, v)
+ in my_columns.iteritems()
+ if not columns or k in columns)
+
+ def multiget(self, keys, columns=None, include_timestamp=False, **kwargs):
+ """Get multiple key values from the column family stub."""
+
+ return OrderedDict(
+ (key, self.get(
+ key,
+ columns,
+ include_timestamp,
+ )) for key in keys if key in self.rows)
+
+ def batch(self, **kwargs):
+ """Returns itself."""
+ return self
+
+ def send(self):
+ pass
+
+ def insert(self, key, columns, timestamp=None, **kwargs):
+ """Insert data to the column family stub."""
+
+ if key not in self.rows:
+ self.rows[key] = OrderedDictWithTime([], timestamp=timestamp)
+
+ for column in columns:
+ self.rows[key].__setitem__(column, columns[column], timestamp)
+
+ return self.rows[key][columns.keys()[0]][1]
+
+ def get_indexed_slices(self, index_clause, **kwargs):
+ """Grabs rows that match a pycassa index clause.
+
+ See :meth:`pycassa.index.create_index_clause()` for creating such an
+ index clause."""
+
+ keys = []
+ for key, row in self.rows.iteritems():
+ for expr in index_clause.expressions:
+ if (
+ expr.column_name in row and
+ operator_dict[expr.op](row[expr.column_name][0], expr.value)
+ ):
+ keys.append(key)
+
+ data = self.multiget(keys, **kwargs).items()
+ return data
+
+ def remove(self, key, columns=None):
+ """Remove a key from the column family stub."""
+ if key not in self.rows:
+ raise NotFoundException()
+ if columns is None:
+ del self.rows[key]
+ else:
+ for c in columns:
+ if c in self.rows[key]:
+ del self.rows[key][c]
+ if not self.rows[key]:
+ del self.rows[key]
+
+
+ def get_range(self, include_timestamp=False, columns=None, **kwargs):
+ """Currently just gets all values from the column family."""
+
+ return [(key, self.get(key, columns, include_timestamp))
+ for key in self.rows]
+
+ def truncate(self):
+ """Clears all data from the column family stub."""
+
+ self.rows.clear()
No changes.
View
@@ -0,0 +1,123 @@
+import unittest
+
+from nose.tools import assert_raises, assert_equal, assert_true
+
+from pycassa import index, ColumnFamily, ConnectionPool,\
+ NotFoundException, SystemManager
+
+from pycassa.contrib.stubs import ColumnFamilyStub, ConnectionPoolStub, \
+ SystemManagerStub
+
+pool = cf = None
+pool_stub = cf_stub = None
+
+
+def setup_module():
+ global pool, cf, indexed_cf, pool_stub, indexed_cf_stub, cf_stub
+ credentials = {'username': 'jsmith', 'password': 'havebadpass'}
+ pool = ConnectionPool(keyspace='PycassaTestKeyspace',
+ credentials=credentials, timeout=1.0)
+ cf = ColumnFamily(pool, 'Standard1', dict_class=TestDict)
+ indexed_cf = ColumnFamily(pool, 'Indexed1')
+
+ pool_stub = ConnectionPoolStub(keyspace='PycassaTestKeyspace',
+ credentials=credentials, timeout=1.0)
+ cf_stub = ColumnFamilyStub(pool_stub, 'Standard1', dict_class=TestDict)
+ indexed_cf_stub = ColumnFamilyStub(pool_stub, 'Indexed1')
+
+
+def teardown_module():
+ cf.truncate()
+ cf_stub.truncate()
+ indexed_cf.truncate()
+ indexed_cf_stub.truncate()
+ pool.dispose()
+
+
+class TestDict(dict):
+ pass
+
+
+class TestColumnFamilyStub(unittest.TestCase):
+
+ def setUp(self):
+ pass
+
+ def tearDown(self):
+ for test_cf in (cf, cf_stub):
+ for key, columns in test_cf.get_range():
+ test_cf.remove(key)
+
+ def test_empty(self):
+ key = 'TestColumnFamily.test_empty'
+
+ for test_cf in (cf, cf_stub):
+ assert_raises(NotFoundException, test_cf.get, key)
+ assert_equal(len(test_cf.multiget([key])), 0)
+ for key, columns in test_cf.get_range():
+ assert_equal(len(columns), 0)
+
+ def test_insert_get(self):
+ key = 'TestColumnFamily.test_insert_get'
+ columns = {'1': 'val1', '2': 'val2'}
+ for test_cf in (cf, cf_stub):
+ assert_raises(NotFoundException, test_cf.get, key)
+ ts = test_cf.insert(key, columns)
+ assert_true(isinstance(ts, (int, long)))
+ assert_equal(test_cf.get(key), columns)
+
+
+ def test_insert_multiget(self):
+ key1 = 'TestColumnFamily.test_insert_multiget1'
+ columns1 = {'1': 'val1', '2': 'val2'}
+ key2 = 'test_insert_multiget1'
+ columns2 = {'3': 'val1', '4': 'val2'}
+ missing_key = 'key3'
+
+ for test_cf in (cf, cf_stub):
+ test_cf.insert(key1, columns1)
+ test_cf.insert(key2, columns2)
+ rows = test_cf.multiget([key1, key2, missing_key])
+ assert_equal(len(rows), 2)
+ assert_equal(rows[key1], columns1)
+ assert_equal(rows[key2], columns2)
+ assert_true(missing_key not in rows)
+
+
+ def insert_insert_get_indexed_slices(self):
+ columns = {'birthdate': 1L}
+
+ keys = set()
+ for i in range(1, 4):
+ indexed_cf.insert('key%d' % i, columns)
+ indexed_cf_stub.insert('key%d' % i, columns)
+ keys.add('key%d' % i)
+
+ expr = index.create_index_expression(column_name='birthdate', value=1L)
+ clause = index.create_index_clause([expr])
+
+ for test_indexed_cf in (indexed_cf, indexed_cf_stub):
+ count = 0
+ for key, cols in test_indexed_cf.get_indexed_slices(clause):
+ assert_equal(cols, columns)
+ assert key in keys
+ count += 1
+ assert_equal(count, 3)
+
+
+ def test_remove(self):
+ key = 'TestColumnFamily.test_remove'
+ for test_cf in (cf, cf_stub):
+ columns = {'1': 'val1', '2': 'val2'}
+ test_cf.insert(key, columns)
+
+ # An empty list for columns shouldn't delete anything
+ test_cf.remove(key, columns=[])
+ assert_equal(test_cf.get(key), columns)
+
+ test_cf.remove(key, columns=['2'])
+ del columns['2']
+ assert_equal(test_cf.get(key), {'1': 'val1'})
+
+ test_cf.remove(key)
+ assert_raises(NotFoundException, test_cf.get, key)

0 comments on commit 957feee

Please sign in to comment.