Update tests to match API changes

bshanks · Jul 17, 2012 · d39f23b · d39f23b
1 parent e121826
commit d39f23b
Show file tree

Hide file tree

Showing 8 changed files with 139 additions and 180 deletions.
diff --git a/pycassa/connection.py b/pycassa/connection.py
@@ -2,7 +2,7 @@
 from thrift.transport import TSocket
 from thrift.protocol import TBinaryProtocol
 
-from pycassa.cassandra.c11 import Cassandra
+from pycassa.cassandra import Cassandra
 from pycassa.cassandra.ttypes import AuthenticationRequest
 
 DEFAULT_SERVER = 'localhost:9160'

diff --git a/pycassa/system_manager.py b/pycassa/system_manager.py
@@ -255,6 +255,7 @@ def create_column_family(self, keyspace, name, column_validation_classes=None, *
             cf_kwargs.setdefault('column_type', 'Super')
 
         for k, v in cf_kwargs.iteritems():
+            v = self._convert_class_attrs(k, v)
             setattr(cfdef, k, v)
 
         if column_validation_classes:
@@ -277,6 +278,7 @@ def alter_column_family(self, keyspace, column_family, column_validation_classes
         cfdef = self.get_keyspace_column_families(keyspace)[column_family]
 
         for k, v in cf_kwargs.iteritems():
+            v = self._convert_class_attrs(k, v)
             setattr(cfdef, k, v)
 
         if column_validation_classes:
@@ -293,6 +295,32 @@ def drop_column_family(self, keyspace, column_family):
         self._conn.set_keyspace(keyspace)
         self._schema_update(self._conn.system_drop_column_family, column_family)
 
+    def _convert_class_attrs(self, attr, value):
+        if attr in ('comparator_type', 'subcomparator_type',
+                    'key_validation_class', 'default_validation_class'):
+            return self._qualify_type_class(value)
+        else:
+            return value
+
+    def _qualify_type_class(self, classname):
+        if classname:
+            if isinstance(classname, types.CassandraType):
+                s = str(classname)
+            elif isinstance(classname, basestring):
+                s = classname
+            else:
+                raise TypeError(
+                        "Column family validators and comparators " \
+                        "must be specified as instances of " \
+                        "pycassa.types.CassandraType subclasses or strings.")
+
+            if s.find('.') == -1:
+                return 'org.apache.cassandra.db.marshal.%s' % s
+            else:
+                return s
+        else:
+            return None
+
     def _alter_column_cfdef(self, cfdef, column, value_type):
         if cfdef.column_type == 'Super':
             packer = marshal.packer_for(cfdef.subcomparator_type)
@@ -301,6 +329,7 @@ def _alter_column_cfdef(self, cfdef, column, value_type):
 
         packed_column = packer(column)
 
+        value_type = self._qualify_type_class(value_type)
         cfdef.column_metadata = cfdef.column_metadata or []
         matched = False
         for c in cfdef.column_metadata:
@@ -367,6 +396,7 @@ def create_index(self, keyspace, column_family, column, value_type,
         packer = marshal.packer_for(cfdef.comparator_type)
         packed_column = packer(column)
 
+        value_type = self._qualify_type_class(value_type)
         coldef = ColumnDef(packed_column, value_type, index_type, index_name)
 
         for c in cfdef.column_metadata:

diff --git a/tests/__init__.py b/tests/__init__.py
@@ -1,5 +1,4 @@
-from pycassa.system_manager import *
-from pycassa.cassandra.constants import *
+from pycassa.system_manager import SystemManager
 
 TEST_KS = 'PycassaTestKeyspace'
 
@@ -8,17 +7,17 @@ def setup_package():
     if TEST_KS in sys.list_keyspaces():
         sys.drop_keyspace(TEST_KS)
     try:
-        sys.create_keyspace(TEST_KS, SIMPLE_STRATEGY, {'replication_factor': '1'})
+        sys.create_keyspace(TEST_KS, 'SimpleStrategy', {'replication_factor': '1'})
         sys.create_column_family(TEST_KS, 'Standard1')
         sys.create_column_family(TEST_KS, 'Super1', super=True)
         sys.create_column_family(TEST_KS, 'Indexed1')
-        sys.create_index(TEST_KS, 'Indexed1', 'birthdate', LONG_TYPE)
-        if sys._conn.version != CASSANDRA_07:
-            sys.create_column_family(TEST_KS, 'Counter1',
-                                     default_validation_class=COUNTER_COLUMN_TYPE)
-            sys.create_column_family(TEST_KS, 'SuperCounter1', super=True,
-                                     default_validation_class=COUNTER_COLUMN_TYPE)
+        sys.create_index(TEST_KS, 'Indexed1', 'birthdate', 'LongType')
+        sys.create_column_family(TEST_KS, 'Counter1',
+                                 default_validation_class='CounterColumnType')
+        sys.create_column_family(TEST_KS, 'SuperCounter1', super=True,
+                                 default_validation_class='CounterColumnType')
     except Exception, e:
+        print e
         try:
             sys.drop_keyspace(TEST_KS)
         except: