consistency with batch verified

cqlengine/query.py

@@ -208,17 +208,22 @@ class BatchQuery(object):
 
     http://www.datastax.com/docs/1.2/cql_cli/cql/BATCH
     """
+    _consistency = None
 
-    def __init__(self, batch_type=None, timestamp=None):
+    def __init__(self, batch_type=None, timestamp=None, consistency=None):
         self.queries = []
         self.batch_type = batch_type
         if timestamp is not None and not isinstance(timestamp, datetime):
             raise CQLEngineException('timestamp object must be an instance of datetime')
         self.timestamp = timestamp
+        self._consistency = consistency
 
     def add_query(self, query, params):
         self.queries.append((query, params))
 
+    def consistency(self, consistency):
+        self._consistency = consistency
+
     def execute(self):
         if len(self.queries) == 0:
             # Empty batch is a no-op
@@ -238,7 +243,7 @@ class BatchQuery(object):
 
         query_list.append('APPLY BATCH;')
 
-        execute('\n'.join(query_list), parameters)
+        execute('\n'.join(query_list), parameters, self._consistency)
 
         self.queries = []
 

cqlengine/tests/test_consistency.py

@@ -5,7 +5,7 @@ from uuid import uuid4
 from cqlengine import columns
 import mock
 from cqlengine.connection import ConnectionPool
-from cqlengine import ALL
+from cqlengine import ALL, BatchQuery
 
 class TestConsistencyModel(Model):
     id      = columns.UUID(primary_key=True, default=lambda:uuid4())
@@ -50,3 +50,18 @@ class TestConsistency(BaseConsistencyTest):
         self.assertEqual(ALL, args[0][2])
 
 
+    def test_batch_consistency(self):
+
+        with mock.patch.object(ConnectionPool, 'execute') as m:
+            with BatchQuery(consistency=ALL) as b:
+                TestConsistencyModel.batch(b).create(text="monkey")
+
+        args = m.call_args
+        self.assertEqual(ALL, args[0][2])
+
+        with mock.patch.object(ConnectionPool, 'execute') as m:
+            with BatchQuery() as b:
+                TestConsistencyModel.batch(b).create(text="monkey")
+
+        args = m.call_args
+        self.assertNotEqual(ALL, args[0][2])