def test_bigquery_magic_w_maximum_bytes_billed_w_context_setter():
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
magics.context._project = None
magics.context.default_query_job_config = job.QueryJobConfig(
maximum_bytes_billed=10203
)
project = "test-project"
job_reference = copy.deepcopy(JOB_REFERENCE_RESOURCE)
job_reference["projectId"] = project
query = "SELECT 17 AS num"
resource = copy.deepcopy(QUERY_RESOURCE)
resource["jobReference"] = job_reference
resource["configuration"]["query"]["query"] = query
data = {"jobReference": job_reference, "totalRows": 0, "rows": []}
credentials_mock = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)
default_patch = mock.patch(
"google.auth.default", return_value=(credentials_mock, "general-project")
)
conn = magics.context._connection = make_connection(resource, data)
list_rows_patch = mock.patch(
"google.cloud.bigquery.client.Client.list_rows",
return_value=google.cloud.bigquery.table._EmptyRowIterator(),
)
with list_rows_patch, default_patch:
ip.run_cell_magic("bigquery", "", query)
_, req = conn.api_request.call_args_list[0]
sent_config = req["data"]["configuration"]["query"]
assert sent_config["maximumBytesBilled"] == "10203"
def test_context_no_connection():
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
magics.context._project = None
magics.context._credentials = None
magics.context._connection = None
credentials_mock = mock.create_autospec(
google.auth.credentials.Credentials, instance=True)
project = "project-123"
default_patch = mock.patch("google.auth.default",
return_value=(credentials_mock, project))
query = "select * from persons"
job_reference = {"projectId": project, "jobId": "some-random-id"}
table = {"projectId": project, "datasetId": "ds", "tableId": "persons"}
resource = {
"jobReference": job_reference,
"configuration": {
"query": {
"destinationTable": table,
"query": query,
"queryParameters": [],
"useLegacySql": False,
}
},
"status": {
"state": "DONE"
},
}
data = {"jobReference": job_reference, "totalRows": 0, "rows": []}
conn_mock = make_connection(resource, data, data, data)
conn_patch = mock.patch("google.cloud.bigquery.client.Connection",
autospec=True)
list_rows_patch = mock.patch(
"google.cloud.bigquery.client.Client.list_rows",
return_value=google.cloud.bigquery.table._EmptyRowIterator(),
)
with conn_patch as conn, list_rows_patch as list_rows, default_patch:
conn.return_value = conn_mock
ip.run_cell_magic("bigquery", "", query)
# Check that query actually starts the job.
list_rows.assert_called()
assert len(conn_mock.api_request.call_args_list) == 2
_, req = conn_mock.api_request.call_args_list[0]
assert req["method"] == "POST"
assert req["path"] == "/projects/{}/jobs".format(project)
sent = req["data"]
assert isinstance(sent["jobReference"]["jobId"], six.string_types)
sent_config = sent["configuration"]["query"]
assert sent_config["query"] == query
def test_context_no_connection():
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
magics.context._project = None
magics.context._credentials = None
magics.context._connection = None
credentials_mock = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)
project = "project-123"
default_patch = mock.patch(
"google.auth.default", return_value=(credentials_mock, project)
)
query = "select * from persons"
job_reference = {"projectId": project, "jobId": "some-random-id"}
table = {"projectId": project, "datasetId": "ds", "tableId": "persons"}
resource = {
"jobReference": job_reference,
"configuration": {
"query": {
"destinationTable": table,
"query": query,
"queryParameters": [],
"useLegacySql": False,
}
},
"status": {"state": "DONE"},
}
data = {"jobReference": job_reference, "totalRows": 0, "rows": []}
conn_mock = make_connection(resource, data, data, data)
conn_patch = mock.patch("google.cloud.bigquery.client.Connection", autospec=True)
list_rows_patch = mock.patch(
"google.cloud.bigquery.client.Client.list_rows",
return_value=google.cloud.bigquery.table._EmptyRowIterator(),
)
with conn_patch as conn, list_rows_patch as list_rows, default_patch:
conn.return_value = conn_mock
ip.run_cell_magic("bigquery", "", query)
# Check that query actually starts the job.
list_rows.assert_called()
assert len(conn_mock.api_request.call_args_list) == 2
_, req = conn_mock.api_request.call_args_list[0]
assert req["method"] == "POST"
assert req["path"] == "/projects/{}/jobs".format(project)
sent = req["data"]
assert isinstance(sent["jobReference"]["jobId"], six.string_types)
sent_config = sent["configuration"]["query"]
assert sent_config["query"] == query
def test_context_no_connection():
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
magics.context._project = None
magics.context._credentials = None
magics.context._connection = None
credentials_mock = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)
project = "project-123"
default_patch = mock.patch(
"google.auth.default", return_value=(credentials_mock, project)
)
job_reference = copy.deepcopy(JOB_REFERENCE_RESOURCE)
job_reference["projectId"] = project
query = "select * from persons"
resource = copy.deepcopy(QUERY_RESOURCE)
resource["jobReference"] = job_reference
resource["configuration"]["query"]["query"] = query
data = {"jobReference": job_reference, "totalRows": 0, "rows": []}
conn_mock = make_connection(resource, data, data, data)
conn_patch = mock.patch("google.cloud.bigquery.client.Connection", autospec=True)
list_rows_patch = mock.patch(
"google.cloud.bigquery.client.Client.list_rows",
return_value=google.cloud.bigquery.table._EmptyRowIterator(),
)
with conn_patch as conn, list_rows_patch as list_rows, default_patch:
conn.return_value = conn_mock
ip.run_cell_magic("bigquery", "", query)
# Check that query actually starts the job.
list_rows.assert_called()
assert len(conn_mock.api_request.call_args_list) == 2
_, req = conn_mock.api_request.call_args_list[0]
assert req["method"] == "POST"
assert req["path"] == "/projects/{}/jobs".format(project)
sent = req["data"]
assert isinstance(sent["jobReference"]["jobId"], six.string_types)
sent_config = sent["configuration"]["query"]
assert sent_config["query"] == query
