def test_to_frame(dummy_db_credentials):
with patch("google.cloud.bigquery.client.Client") as connection, patch(
"google.cloud.bigquery.client.Client.query") as bq_query:
bq_cli = BigQueryClient(**dummy_db_credentials)
bq_cli.credentials = True
bq_cli.connection = connection
assert bq_query.call_count == 0
_ = bq_cli.to_frame("SELECT *")
assert bq_query.call_count == 1
def test_insert_from_frame_create_first_true_with_connection(
dummy_db_credentials, df):
with patch("google.cloud.bigquery.client.Client") as connection:
bq_cli = BigQueryClient(**dummy_db_credentials)
bq_cli.credentials = True
bq_cli.insert_from_frame(df, create_first=True, connection=connection)
connection.load_table_from_dataframe.assert_called_once_with(
df, bq_cli.table_id)
connection.load_table_from_dataframe.return_value.result.assert_called_once(
)
def test_bigquery_close(dummy_db_credentials):
with patch("google.cloud.bigquery.client.Client") as connection, patch(
"google.cloud.bigquery.client.Client.close") as bq_close:
bq_cli = BigQueryClient(**dummy_db_credentials)
bq_cli.credentials = True
bq_cli.connection = connection
bq_cli.execute("SELECT *")
assert bq_close.call_count == 1
bq_cli.close()
assert bq_close.call_count == 2
def test_insert_from_frame_create_first_true_get_table_not_found(
dummy_db_credentials, df):
from google.cloud import exceptions
bq_cli = BigQueryClient(**dummy_db_credentials)
bq_cli.credentials = True
with patch(
"google.cloud.bigquery.client.Client") as connection, patch.object(
bq_cli, 'execute') as execute_mock:
bq_cli.connection = connection
connection.get_table.side_effect = exceptions.NotFound('Not Found')
q = "CREATE TABLE {table_id}"
params = {'table_id': 'test'}
bq_cli.insert_from_frame(
df,
table=params['table_id'],
create_first=True,
create_sql=q.format(**params),
connection=connection,
)
connection.get_table.assert_called_once_with(params['table_id'])
execute_mock.assert_called_once_with(
q.format(**params),
params=params,
connection=connection,
)
def test_insert_from_frame(dummy_db_credentials, df):
with patch("google.cloud.bigquery.client.Client") as connection, patch(
"google.cloud.bigquery.client.Client.load_table_from_dataframe"
) as bq_load:
bq_cli = BigQueryClient(**dummy_db_credentials)
bq_cli.credentials = True
bq_cli.connection = connection
assert bq_load.call_count == 0
bq_cli.insert_from_frame(df, create_first=False)
assert bq_load.call_count == 1
def test_connect_credentials_none(dummy_db_credentials):
with patch("google.cloud.bigquery.client.Client") as connection, patch(
"google.cloud.bigquery.client.Client.query") as bq_query, patch(
"muttlib.dbconn.bigquery.BigQueryClient._read_cred"
) as mock_read_cred:
bq_cli = BigQueryClient(**dummy_db_credentials)
bq_cli.credentials = None
bq_cli.connection = connection
bq_cli.execute("SELECT *")
mock_read_cred.assert_called_once_with(bq_cli.auth, bq_cli.auth_file)
def test_insert_from_frame_create_first_true(dummy_db_credentials):
with patch("google.cloud.bigquery.client.Client") as connection, patch(
"google.cloud.bigquery.client.Client.load_table_from_dataframe"
) as bq_load:
df = pd.DataFrame({'col1': [1], 'col2': [3.0]})
bq_cli = BigQueryClient(**dummy_db_credentials)
bq_cli.credentials = True
bq_cli.connection = connection
bq_cli.insert_from_frame(df, create_first=True)
bq_load.assert_called_once_with(df, bq_cli.table_id)
def test_execute_bigquery_with_params(dummy_db_credentials):
with patch("google.cloud.bigquery.client.Client") as connection, patch(
"google.cloud.bigquery.client.Client.query") as bq_query:
bq_cli = BigQueryClient(**dummy_db_credentials)
bq_cli.credentials = True
bq_cli.connection = connection
q = """
SELECT * FROM {table}
WHERE {condition1}
"""
params = {"table": "teradactyl", "condition1": "id = 1"}
bq_cli.execute(q, connection=connection, params=params)
connection.query.assert_called_once_with(q.format(**params))
def test_execute_bigquery_with_incomplete_params_fails(dummy_db_credentials):
with patch("google.cloud.bigquery.client.Client") as connection, patch(
"google.cloud.bigquery.client.Client.query") as bq_query:
bq_cli = BigQueryClient(**dummy_db_credentials)
bq_cli.credentials = True
bq_cli.connection = connection
q = """
SELECT * FROM {table}
WHERE {condition1}
"""
with pytest.raises(KeyError):
bq_cli.execute(q,
connection=connection,
params={"table": "teradactyl"})