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)