def test_to_dataframe_none(mock_pandas, dbapi, rs_creds):
with mock.patch(dbapi.__name__ + '.connect') as mock_connect:
mock_connect.return_value.cursor.return_value.fetchmany.return_value = []
with locopy.Cmd(dbapi=dbapi, **rs_creds) as cmd:
cmd.execute("SELECT 'hello world' WHERE 1=0")
assert cmd.to_dataframe(size=5) is None
mock_pandas.assert_not_called()
def test_redshift_connect(rs_creds, dbapi):
with mock.patch(dbapi.__name__ + '.connect') as mock_connect:
test_redshift = locopy.Cmd(dbapi=dbapi, **rs_creds)
if dbapi.__name__ == 'pg8000':
mock_connect.assert_called_with(host='host',
user='******',
port='port',
password='******',
database='dbname',
ssl=True)
else:
mock_connect.assert_called_with(host='host',
user='******',
port='port',
password='******',
database='dbname',
sslmode='require')
test_redshift.conn.cursor.assert_called_with()
# side effect exception
mock_connect.side_effect = Exception('Connect Exception')
with pytest.raises(RedshiftConnectionError):
test_redshift._connect()
def test_cmd_execute_single_rows(dbapi):
expected = pd.DataFrame({'field_1': [1], 'field_2': [2]})
with locopy.Cmd(dbapi=dbapi, **CREDS_DICT) as cmd:
cmd.execute("SELECT 1 AS field_1, 2 AS field_2 ")
df = cmd.to_dataframe()
assert np.allclose(df['field_1'], expected['field_1'])
def test_execute_no_connection_exception(rs_creds, dbapi):
with mock.patch(dbapi.__name__ + '.connect') as mock_connect:
test = locopy.Cmd(dbapi=dbapi, **rs_creds)
test.conn = None
test.cursor = None
with pytest.raises(RedshiftConnectionError):
test.execute('SELECT * FROM some_table')
def test_to_dataframe_all(mock_pandas, dbapi, rs_creds):
with mock.patch(dbapi.__name__ + '.connect') as mock_connect:
mock_connect.return_value.cursor.return_value.fetchall.return_value = [
(1, 2), (2, 3), (3, )
]
with locopy.Cmd(dbapi=dbapi, **rs_creds) as cmd:
cmd.execute("SELECT 'hello world' AS fld")
df = cmd.to_dataframe()
assert mock_connect.return_value.cursor.return_value.fetchall.called
mock_pandas.assert_called_with(cmd.cursor.fetchall(), columns=[])
def test_redshift_disconnect(rs_creds, dbapi):
with mock.patch(dbapi.__name__ + '.connect') as mock_connect:
test_redshift = locopy.Cmd(dbapi=dbapi, **rs_creds)
test_redshift._disconnect()
test_redshift.conn.close.assert_called_with()
test_redshift.cursor.close.assert_called_with()
# side effect exception
test_redshift.conn.close.side_effect = Exception('Disconnect Exception')
with pytest.raises(RedshiftDisconnectionError):
test_redshift._disconnect()
def test_is_connected(rs_creds, dbapi):
with mock.patch(dbapi.__name__ + '.connect') as mock_connect:
test_redshift = CmdNoConnect(dbapi=dbapi, **rs_creds)
assert test_redshift._is_connected() is False
with mock.patch(dbapi.__name__ + '.connect') as mock_connect:
test_redshift = locopy.Cmd(dbapi=dbapi, **rs_creds)
assert test_redshift._is_connected() is True
# throws exception in _is_connected
test_redshift = CmdNoConnect(**rs_creds)
del test_redshift.conn
assert test_redshift._is_connected() is False
def test_get_redshift(rs_creds, dbapi):
with mock.patch(dbapi.__name__ + '.connect') as mock_connect:
locopy.Cmd(dbapi=dbapi, **rs_creds)
if dbapi.__name__ == 'pg8000':
mock_connect.assert_called_with(host='host',
user='******',
port='port',
password='******',
database='dbname',
ssl=True)
else:
mock_connect.assert_called_with(host='host',
user='******',
port='port',
password='******',
database='dbname',
sslmode='require')
def test_execute_sql_exception(rs_creds, dbapi):
with mock.patch(dbapi.__name__ + '.connect') as mock_connect:
with locopy.Cmd(dbapi=dbapi, **rs_creds) as test:
test.cursor.execute.side_effect = Exception('SQL Exception')
with pytest.raises(RedshiftError):
test.execute('SELECT * FROM some_table')
def test_execute(rs_creds, dbapi):
with mock.patch(dbapi.__name__ + '.connect') as mock_connect:
with locopy.Cmd(dbapi=dbapi, **rs_creds) as test:
test.execute('SELECT * FROM some_table')
assert test.cursor.execute.called is True
