def test_connect_with_pkey(self, mock_serialization): mock_snowflake = Mock(name='mock_snowflake') mock_connector = mock_snowflake.connector mock_pkey = mock_serialization.load_pem_private_key.return_value = Mock(name='pkey') # need to patch this here so it can be imported in the function scope with patch.dict('sys.modules', snowflake=mock_snowflake): mock_connector.connect.return_value = 'OK' snowflake = SnowflakeDatabase(user='******', private_key_data='abcdefg', private_key_password='******', account='test_account', database='test_database') result = snowflake.connect() with self.subTest('returns connection'): self.assertEqual('OK', result) with self.subTest('connects with credentials'): mock_serialization.load_pem_private_key.assert_called_once_with(b'abcdefg', b'1234', backend=ANY) with self.subTest('connects with credentials'): mock_connector.connect.assert_called_once_with(user='******', password=None, account='test_account', database='test_database', private_key=mock_pkey.private_bytes.return_value, region=None, warehouse=None)
def test_get_column_definitions(self, mock_fetch): SnowflakeDatabase().get_column_definitions('test_schema', 'test_table') mock_fetch.assert_called_once_with( 'SELECT DISTINCT COLUMN_NAME,DATA_TYPE ' 'FROM INFORMATION_SCHEMA.COLUMNS ' 'WHERE TABLE_SCHEMA=%(schema)s AND TABLE_NAME=%(table)s ' 'ORDER BY column_name', connection=None, parameters={'schema': 'test_schema', 'table': 'test_table'} )
def test_defaults(self): snowflake = SnowflakeDatabase() self.assertEqual('snowflake', snowflake.account) self.assertEqual('snowflake', snowflake.database) self.assertEqual('snowflake', snowflake.user) self.assertIsNone(snowflake.password) self.assertIsNone(snowflake.private_key_data) self.assertIsNone(snowflake.private_key_password) self.assertIsNone(snowflake.region) self.assertIsNone(snowflake.warehouse)
def test_connect_with_password(self): mock_snowflake = Mock(name='mock_snowflake') mock_connector = mock_snowflake.connector # need to patch this here so it can be imported in the function scope with patch.dict('sys.modules', snowflake=mock_snowflake): mock_connector.connect.return_value = 'OK' snowflake = SnowflakeDatabase(user='******', password='******', account='test_account', database='test_database') result = snowflake.connect() self.assertEqual('OK', result) mock_connector.connect.assert_called_once_with(user='******', password='******', account='test_account', database='test_database', private_key=None, region=None, warehouse=None)
def test_trunc_quarter(self): result = SnowflakeDatabase().trunc_date(Field('date'), 'quarter') self.assertEqual('TRUNC("date",\'Q\')', str(result))
def test_trunc_week(self): result = SnowflakeDatabase().trunc_date(Field('date'), 'week') self.assertEqual('TRUNC("date",\'IW\')', str(result))
def test_date_add_year(self): result = SnowflakeDatabase().date_add(Field('date'), 'year', 1) self.assertEqual('TIMESTAMPADD(\'year\',1,"date")', str(result))
def test_trunc_day(self): result = SnowflakeDatabase().trunc_date(Field('date'), daily) self.assertEqual('TRUNC("date",\'DD\')', str(result))
def test_trunc_year(self): result = SnowflakeDatabase().trunc_date(Field('date'), annually) self.assertEqual('TRUNC("date",\'Y\')', str(result))
def test_get_column_definitions(self, mock_fetch): SnowflakeDatabase().get_column_definitions('test_schema', 'test_table') mock_fetch.assert_called_once_with( 'DESCRIBE TABLE test_schema.test_table TYPE=COLUMNS', connection=None)
def test_date_add_month(self): result = SnowflakeDatabase().date_add(Field('date'), 'month', 1) self.assertEqual('TIMESTAMPADD(month,1,"date")', str(result))
def test_date_add_week(self): result = SnowflakeDatabase().date_add(Field('date'), 'week', 1) self.assertEqual('TIMESTAMPADD(week,1,"date")', str(result))