def setUp(self): """Set up the test by creating a query with multiple batches""" self.statement_list = statement_list = [ 'select version;', 'select * from t1;' ] self.statement_str = ''.join(statement_list) self.query_uri = 'test_uri' self.query = Query( self.query_uri, self.statement_str, QueryExecutionSettings(ExecutionPlanOptions(), ResultSetStorageType.FILE_STORAGE), QueryEvents()) self.mock_query_results = [('Id1', 'Value1'), ('Id2', 'Value2')] self.cursor = MockCursor(self.mock_query_results) self.connection = MockPGServerConnection(cur=self.cursor) self.columns_info = [] db_column_id = DbColumn() db_column_id.data_type = 'text' db_column_id.column_name = 'Id' db_column_id.provider = PG_PROVIDER_NAME db_column_value = DbColumn() db_column_value.data_type = 'text' db_column_value.column_name = 'Value' db_column_value.provider = PG_PROVIDER_NAME self.columns_info = [db_column_id, db_column_value] self.get_columns_info_mock = mock.Mock(return_value=self.columns_info)
def test_write_tstzrange(self): test_value = DateTimeTZRange( datetime.datetime(2014, 6, 8, 12, 12, 45, tzinfo=psycopg2.tz.FixedOffsetTimezone( offset=720, name=None)), datetime.datetime(2016, 7, 6, 14, 12, 8, tzinfo=psycopg2.tz.FixedOffsetTimezone( offset=720, name=None))) test_columns_info = [] col = DbColumn() col.data_type = datatypes.DATATYPE_TSTZRANGE col.provider = PG_PROVIDER_NAME test_columns_info.append(col) mock_storage_data_reader = MockStorageDataReader( self._cursor, test_columns_info) mock_storage_data_reader.get_value = mock.MagicMock( return_value=test_value) res = self._writer.write_row(mock_storage_data_reader) self.assertEqual( self.get_expected_length_with_additional_buffer_for_size( len("[2014-06-08T12:12:45+12:00,2016-07-06T14:12:08+12:00)")), res)
def test_write_short(self): test_value = 12345 test_columns_info = [] col = DbColumn() col.data_type = datatypes.DATATYPE_SMALLINT col.provider = PG_PROVIDER_NAME test_columns_info.append(col) mock_storage_data_reader = MockStorageDataReader( self._cursor, test_columns_info) mock_storage_data_reader.get_value = mock.MagicMock( return_value=test_value) res = self._writer.write_row(mock_storage_data_reader) self.assertEqual( self.get_expected_length_with_additional_buffer_for_size(2), res)
def test_write_udt(self): test_value = "TestUserDefinedTypes" test_columns_info = [] col = DbColumn() col.data_type = 'UserDefinedTypes' col.provider = PG_PROVIDER_NAME test_columns_info.append(col) mock_storage_data_reader = MockStorageDataReader( self._cursor, test_columns_info) mock_storage_data_reader.get_value = mock.MagicMock( return_value=test_value) res = self._writer.write_row(mock_storage_data_reader) self.assertEqual( self.get_expected_length_with_additional_buffer_for_size( len(test_value)), res)
def test_write_int4range(self): test_value = NumericRange(10, 20) test_columns_info = [] col = DbColumn() col.data_type = datatypes.DATATYPE_INT4RANGE col.provider = PG_PROVIDER_NAME test_columns_info.append(col) mock_storage_data_reader = MockStorageDataReader( self._cursor, test_columns_info) mock_storage_data_reader.get_value = mock.MagicMock( return_value=test_value) res = self._writer.write_row(mock_storage_data_reader) self.assertEqual( self.get_expected_length_with_additional_buffer_for_size( len("[10,20)")), res)
def test_write_json(self): test_value = {"Name": "TestName", "Schema": "TestSchema"} test_columns_info = [] col = DbColumn() col.data_type = datatypes.DATATYPE_JSON col.provider = PG_PROVIDER_NAME test_columns_info.append(col) mock_storage_data_reader = MockStorageDataReader( self._cursor, test_columns_info) mock_storage_data_reader.get_value = mock.MagicMock( return_value=test_value) res = self._writer.write_row(mock_storage_data_reader) self.assertEqual( self.get_expected_length_with_additional_buffer_for_size( len(str(test_value))), res)
def test_write_bytea(self): test_value = memoryview(b'TestString') test_columns_info = [] col = DbColumn() col.data_type = datatypes.DATATYPE_BYTEA col.provider = PG_PROVIDER_NAME test_columns_info.append(col) mock_storage_data_reader = MockStorageDataReader( self._cursor, test_columns_info) mock_storage_data_reader.get_value = mock.MagicMock( return_value=test_value) res = self._writer.write_row(mock_storage_data_reader) self.assertEqual( self.get_expected_length_with_additional_buffer_for_size( len(test_value.tobytes())), res)
def test_write_uuid(self): test_value = uuid.uuid4() test_columns_info = [] col = DbColumn() col.data_type = datatypes.DATATYPE_UUID col.provider = PG_PROVIDER_NAME test_columns_info.append(col) mock_storage_data_reader = MockStorageDataReader( self._cursor, test_columns_info) mock_storage_data_reader.get_value = mock.MagicMock( return_value=test_value) res = self._writer.write_row(mock_storage_data_reader) self.assertEqual( self.get_expected_length_with_additional_buffer_for_size(36), res) # UUID standard len is 36
def test_write_timedelta(self): test_value = datetime.timedelta(days=3, hours=4, minutes=5, seconds=6) test_columns_info = [] col = DbColumn() col.data_type = datatypes.DATATYPE_INTERVAL col.provider = PG_PROVIDER_NAME test_columns_info.append(col) mock_storage_data_reader = MockStorageDataReader( self._cursor, test_columns_info) mock_storage_data_reader.get_value = mock.MagicMock( return_value=test_value) res = self._writer.write_row(mock_storage_data_reader) self.assertEqual( self.get_expected_length_with_additional_buffer_for_size( len(str(test_value))), res)
def test_write_datetime(self): test_value = datetime.datetime(2004, 10, 19, 10, 23, 54) test_columns_info = [] col = DbColumn() col.data_type = datatypes.DATATYPE_TIMESTAMP col.provider = PG_PROVIDER_NAME test_columns_info.append(col) mock_storage_data_reader = MockStorageDataReader( self._cursor, test_columns_info) mock_storage_data_reader.get_value = mock.MagicMock( return_value=test_value) res = self._writer.write_row(mock_storage_data_reader) self.assertEqual( self.get_expected_length_with_additional_buffer_for_size( len(test_value.isoformat())), res)
def test_write_daterange(self): test_value = DateRange(datetime.date(2015, 6, 6), datetime.date(2016, 8, 8)) test_columns_info = [] col = DbColumn() col.data_type = datatypes.DATATYPE_DATERANGE col.provider = PG_PROVIDER_NAME test_columns_info.append(col) mock_storage_data_reader = MockStorageDataReader( self._cursor, test_columns_info) mock_storage_data_reader.get_value = mock.MagicMock( return_value=test_value) res = self._writer.write_row(mock_storage_data_reader) self.assertEqual( self.get_expected_length_with_additional_buffer_for_size( len("[2015-06-06,2016-08-08)")), res)