def test_datasource(): with pytest.raises(ValidationError): PostgresDataSource(name='mycon', domain='mydomain', database='ubuntu', query='') with pytest.raises(ValueError) as exc_info: PostgresDataSource(name='mycon', domain='mydomain', database='ubuntu') assert "'query' or 'table' must be set" in str(exc_info.value) ds = PostgresDataSource(name='mycon', domain='mydomain', database='ubuntu', table='test') assert ds.query == 'select * from test;'
def test_postgress_get_df(mocker): snock = mocker.patch('psycopg2.connect') reasq = mocker.patch('pandas.read_sql') postgres_connector = PostgresConnector(name='test', host='localhost', user='******', password='******', port=22) ds = PostgresDataSource( domain='test', name='test', database='postgres_db', query='SELECT Name, CountryCode, Population FROM City LIMIT 2;', ) postgres_connector.get_df(ds) snock.assert_called_once_with(host='localhost', dbname='postgres_db', user='******', password='******', port=22) reasq.assert_called_once_with( 'SELECT Name, CountryCode, Population FROM City LIMIT 2;', con=snock(), params={})
def test_open_connection(): """ It should not open a connection """ with pytest.raises(psycopg2.OperationalError): ds = PostgresDataSource(domain='pika', name='pika', database='circle_test', query='q') PostgresConnector(name='test', host='lolcathost', user='******', connect_timeout=1).get_df( ds )
def test_retrieve_response(postgres_connector): """ It should connect to the database and retrieve the response to the query """ ds = PostgresDataSource( domain='test', name='test', query='SELECT Name, CountryCode, Population FROM City LIMIT 2;') res = postgres_connector.get_df(ds) assert isinstance(res, pd.DataFrame) assert res.shape == (2, 3)
def test_get_form_empty_query(postgres_connector): """It should give suggestions of the databases without changing the rest""" current_config = {} form = PostgresDataSource.get_form(postgres_connector, current_config) assert form['properties']['database'] == {'$ref': '#/definitions/database'} assert form['definitions']['database'] == { 'title': 'database', 'description': 'An enumeration.', 'type': 'string', 'enum': ['postgres', 'postgres_db'], }
def test_get_df_array_interpolation(postgres_connector): data_source_spec = { 'domain': 'Postgres test', 'type': 'external_database', 'name': 'Some Postgres provider', 'database': 'postgres_db', 'query': 'SELECT * FROM City WHERE id in %(ids)s', 'parameters': {'ids': [1, 2]}, } data_source = PostgresDataSource(**data_source_spec) df = postgres_connector.get_df(data_source) assert not df.empty assert df.shape == (2, 5)
def test_describe(mocker, postgres_connector): """It should return a table description""" ds = PostgresDataSource( domain='test', name='test', database='postgres_db', query='SELECT * FROM city;' ) res = postgres_connector.describe(ds) assert res == { 'id': 'int4', 'name': 'text', 'countrycode': 'bpchar', 'district': 'text', 'population': 'int4', }
def test_get_df_forbidden_table_interpolation(postgres_connector): data_source_spec = { 'domain': 'Postgres test', 'type': 'external_database', 'name': 'Some Postgres provider', 'database': 'postgres_db', 'query': 'SELECT * FROM %(tablename)s WHERE Population > 5000000', 'parameters': {'tablename': 'City'}, } data_source = PostgresDataSource(**data_source_spec) with pytest.raises(pd.io.sql.DatabaseError) as e: postgres_connector.get_df(data_source) assert 'interpolating table name is forbidden' in str(e.value)
def test_get_df_db_jinja_syntax(postgres_connector): data_source_spec = { 'domain': 'Postgres test', 'type': 'external_database', 'name': 'Some Postgres provider', 'database': 'postgres_db', 'query': 'SELECT * FROM City WHERE Population > {{ max_pop }}', 'parameters': {'max_pop': 5000000}, } expected_columns = {'id', 'name', 'countrycode', 'district', 'population'} data_source = PostgresDataSource(**data_source_spec) df = postgres_connector.get_df(data_source) assert not df.empty assert set(df.columns) == expected_columns assert df.shape == (24, 5)
def test_get_df_db(postgres_connector): """ It should extract the table City and make some merge with some foreign key. """ data_source_spec = { 'domain': 'Postgres test', 'type': 'external_database', 'name': 'Some Postgres provider', 'database': 'postgres_db', 'query': 'SELECT * FROM City WHERE Population > %(max_pop)s', 'parameters': {'max_pop': 5000000}, } expected_columns = {'id', 'name', 'countrycode', 'district', 'population'} data_source = PostgresDataSource(**data_source_spec) df = postgres_connector.get_df(data_source) assert not df.empty assert set(df.columns) == expected_columns assert df.shape == (24, 5)
def test_get_form_query_with_good_database(postgres_connector, mocker): """It should give suggestions of the collections""" current_config = {'database': 'postgres_db'} form = PostgresDataSource.get_form(postgres_connector, current_config) assert form['properties']['database'] == {'$ref': '#/definitions/database'} assert form['definitions']['database'] == { 'title': 'database', 'description': 'An enumeration.', 'type': 'string', 'enum': ['postgres', 'postgres_db'], } assert form['properties']['table'] == {'$ref': '#/definitions/table'} assert form['definitions']['table'] == { 'title': 'table', 'description': 'An enumeration.', 'type': 'string', 'enum': ['city', 'country', 'countrylanguage'], }
def test_describe_error(mocker, postgres_connector): """It should return a table description""" ds = PostgresDataSource( domain='test', name='test', database='postgres_db', query='SELECT * FROM city;' ) mocked_connect = mocker.MagicMock() mocked_cursor = mocker.MagicMock() mocked_connect.cursor.return_value = mocked_cursor def execute(q): raise psycopg2.ProgrammingError mocked_cursor.__enter__().execute = execute mocker.patch( 'toucan_connectors.postgres.postgresql_connector.pgsql.connect', return_value=mocked_connect ) with pytest.raises(psycopg2.ProgrammingError): postgres_connector.describe(ds)
def test_raise_on_empty_query(): with pytest.raises(ValidationError): PostgresDataSource(domaine='test', name='test', database='postgres_db', query='')
def test_get_form_connection_fails(mocker, postgres_connector): """It should return a form even if connect fails""" mocker.patch.object(pgsql, 'connect').side_effect = IOError form = PostgresDataSource.get_form(postgres_connector, current_config={}) assert 'table' in form['properties']