# Copyright 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. import unittest from unittest.mock import Mock from cassandra import ProtocolVersion, UnsupportedOperation from cassandra.protocol import ( PrepareMessage, QueryMessage, ExecuteMessage, UnsupportedOperation, _PAGING_OPTIONS_FLAG, _WITH_SERIAL_CONSISTENCY_FLAG, _PAGE_SIZE_FLAG, _WITH_PAGING_STATE_FLAG, BatchMessage ) from cassandra.query import BatchType from cassandra.marshal import uint32_unpack from cassandra.cluster import ContinuousPagingOptions class MessageTest(unittest.TestCase): def test_prepare_message(self): """ Test to check the appropriate calls are made @since 3.9 @jira_ticket PYTHON-713 @expected_result the values are correctly written @test_category connection """ message = PrepareMessage("a") io = Mock() message.send_body(io, 4) self._check_calls(io, [(b'\x00\x00\x00\x01',), (b'a',)]) io.reset_mock() message.send_body(io, 5) self._check_calls(io, [(b'\x00\x00\x00\x01',), (b'a',), (b'\x00\x00\x00\x00',)]) def test_execute_message(self): message = ExecuteMessage('1', [], 4) io = Mock() message.send_body(io, 4) self._check_calls(io, [(b'\x00\x01',), (b'1',), (b'\x00\x04',), (b'\x01',), (b'\x00\x00',)]) io.reset_mock() message.result_metadata_id = 'foo' message.send_body(io, 5) self._check_calls(io, [(b'\x00\x01',), (b'1',), (b'\x00\x03',), (b'foo',), (b'\x00\x04',), (b'\x00\x00\x00\x01',), (b'\x00\x00',)]) def test_query_message(self): """ Test to check the appropriate calls are made @since 3.9 @jira_ticket PYTHON-713 @expected_result the values are correctly written @test_category connection """ message = QueryMessage("a", 3) io = Mock() message.send_body(io, 4) self._check_calls(io, [(b'\x00\x00\x00\x01',), (b'a',), (b'\x00\x03',), (b'\x00',)]) io.reset_mock() message.send_body(io, 5) self._check_calls(io, [(b'\x00\x00\x00\x01',), (b'a',), (b'\x00\x03',), (b'\x00\x00\x00\x00',)]) def _check_calls(self, io, expected): self.assertEqual( tuple(c[1] for c in io.write.mock_calls), tuple(expected) ) def test_continuous_paging(self): """ Test to check continuous paging throws an Exception if it's not supported and the correct valuesa are written to the buffer if the option is enabled. @since DSE 2.0b3 GRAPH 1.0b1 @jira_ticket PYTHON-694 @expected_result the values are correctly written @test_category connection """ max_pages = 4 max_pages_per_second = 3 continuous_paging_options = ContinuousPagingOptions(max_pages=max_pages, max_pages_per_second=max_pages_per_second) message = QueryMessage("a", 3, continuous_paging_options=continuous_paging_options) io = Mock() for version in [version for version in ProtocolVersion.SUPPORTED_VERSIONS if not ProtocolVersion.has_continuous_paging_support(version)]: self.assertRaises(UnsupportedOperation, message.send_body, io, version) io.reset_mock() message.send_body(io, ProtocolVersion.DSE_V1) # continuous paging adds two write calls to the buffer self.assertEqual(len(io.write.mock_calls), 6) # Check that the appropriate flag is set to True self.assertEqual(uint32_unpack(io.write.mock_calls[3][1][0]) & _WITH_SERIAL_CONSISTENCY_FLAG, 0) self.assertEqual(uint32_unpack(io.write.mock_calls[3][1][0]) & _PAGE_SIZE_FLAG, 0) self.assertEqual(uint32_unpack(io.write.mock_calls[3][1][0]) & _WITH_PAGING_STATE_FLAG, 0) self.assertEqual(uint32_unpack(io.write.mock_calls[3][1][0]) & _PAGING_OPTIONS_FLAG, _PAGING_OPTIONS_FLAG) # Test max_pages and max_pages_per_second are correctly written self.assertEqual(uint32_unpack(io.write.mock_calls[4][1][0]), max_pages) self.assertEqual(uint32_unpack(io.write.mock_calls[5][1][0]), max_pages_per_second) def test_prepare_flag(self): """ Test to check the prepare flag is properly set, This should only happen for V5 at the moment. @since 3.9 @jira_ticket PYTHON-694, PYTHON-713 @expected_result the values are correctly written @test_category connection """ message = PrepareMessage("a") io = Mock() for version in ProtocolVersion.SUPPORTED_VERSIONS: message.send_body(io, version) if ProtocolVersion.uses_prepare_flags(version): self.assertEqual(len(io.write.mock_calls), 3) else: self.assertEqual(len(io.write.mock_calls), 2) io.reset_mock() def test_prepare_flag_with_keyspace(self): message = PrepareMessage("a", keyspace='ks') io = Mock() for version in ProtocolVersion.SUPPORTED_VERSIONS: if ProtocolVersion.uses_keyspace_flag(version): message.send_body(io, version) self._check_calls(io, [ (b'\x00\x00\x00\x01',), (b'a',), (b'\x00\x00\x00\x01',), (b'\x00\x02',), (b'ks',), ]) else: with self.assertRaises(UnsupportedOperation): message.send_body(io, version) io.reset_mock() def test_keyspace_flag_raises_before_v5(self): keyspace_message = QueryMessage('a', consistency_level=3, keyspace='ks') io = Mock(name='io') with self.assertRaisesRegex(UnsupportedOperation, 'Keyspaces.*set'): keyspace_message.send_body(io, protocol_version=4) io.assert_not_called() def test_keyspace_written_with_length(self): io = Mock(name='io') base_expected = [ (b'\x00\x00\x00\x01',), (b'a',), (b'\x00\x03',), (b'\x00\x00\x00\x80',), # options w/ keyspace flag ] QueryMessage('a', consistency_level=3, keyspace='ks').send_body( io, protocol_version=5 ) self._check_calls(io, base_expected + [ (b'\x00\x02',), # length of keyspace string (b'ks',), ]) io.reset_mock() QueryMessage('a', consistency_level=3, keyspace='keyspace').send_body( io, protocol_version=5 ) self._check_calls(io, base_expected + [ (b'\x00\x08',), # length of keyspace string (b'keyspace',), ]) def test_batch_message_with_keyspace(self): self.maxDiff = None io = Mock(name='io') batch = BatchMessage( batch_type=BatchType.LOGGED, queries=((False, 'stmt a', ('param a',)), (False, 'stmt b', ('param b',)), (False, 'stmt c', ('param c',)) ), consistency_level=3, keyspace='ks' ) batch.send_body(io, protocol_version=5) self._check_calls(io, ((b'\x00',), (b'\x00\x03',), (b'\x00',), (b'\x00\x00\x00\x06',), (b'stmt a',), (b'\x00\x01',), (b'\x00\x00\x00\x07',), ('param a',), (b'\x00',), (b'\x00\x00\x00\x06',), (b'stmt b',), (b'\x00\x01',), (b'\x00\x00\x00\x07',), ('param b',), (b'\x00',), (b'\x00\x00\x00\x06',), (b'stmt c',), (b'\x00\x01',), (b'\x00\x00\x00\x07',), ('param c',), (b'\x00\x03',), (b'\x00\x00\x00\x80',), (b'\x00\x02',), (b'ks',)) )