BatchStatement.clear + unit test

aholmberg · aholmberg · commit 4cd4e987f40e · 2016-05-06T15:37:09.000-05:00
diff --git a/cassandra/query.py b/cassandra/query.py
@@ -699,6 +699,19 @@ def __init__(self, batch_type=BatchType.LOGGED, retry_policy=None,
         Statement.__init__(self, retry_policy=retry_policy, consistency_level=consistency_level,
                            serial_consistency_level=serial_consistency_level, custom_payload=custom_payload)
 
+    def clear(self):
+        """
+        This is a convenience method to clear a batch statement for reuse.
+
+        *Note:* it should not be used concurrently with uncompleted execution futures executing the same
+        ``BatchStatement``.
+        """
+        del self._statements_and_parameters[:]
+        self.keyspace = None
+        self.routing_key = None
+        if self.custom_payload:
+            self.custom_payload.clear()
+
     def add(self, statement, parameters=None):
         """
         Adds a :class:`.Statement` and optional sequence of parameters
diff --git a/tests/unit/test_query.py b/tests/unit/test_query.py
@@ -0,0 +1,59 @@
+# Copyright 2013-2016 DataStax, Inc.
+#
+# 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.
+
+try:
+    import unittest2 as unittest
+except ImportError:
+    import unittest  # noqa
+
+import six
+
+from cassandra.query import BatchStatement, SimpleStatement
+
+
+class BatchStatementTest(unittest.TestCase):
+    # TODO: this suite could be expanded; for now just adding a test covering a PR
+
+    def test_clear(self):
+        keyspace = 'keyspace'
+        routing_key = 'routing_key'
+        custom_payload = {'key': six.b('value')}
+
+        ss = SimpleStatement('whatever', keyspace=keyspace, routing_key=routing_key, custom_payload=custom_payload)
+
+        batch = BatchStatement()
+        batch.add(ss)
+
+        self.assertTrue(batch._statements_and_parameters)
+        self.assertEqual(batch.keyspace, keyspace)
+        self.assertEqual(batch.routing_key, routing_key)
+        self.assertEqual(batch.custom_payload, custom_payload)
+
+        batch.clear()
+        self.assertFalse(batch._statements_and_parameters)
+        self.assertIsNone(batch.keyspace)
+        self.assertIsNone(batch.routing_key)
+        self.assertFalse(batch.custom_payload)
+
+        batch.add(ss)
+
+    def test_clear_empty(self):
+        batch = BatchStatement()
+        batch.clear()
+        self.assertFalse(batch._statements_and_parameters)
+        self.assertIsNone(batch.keyspace)
+        self.assertIsNone(batch.routing_key)
+        self.assertFalse(batch.custom_payload)
+
+        batch.add('something')
