From 7b843d072f146099bbbb59b475b3a7cc02a991c5 Mon Sep 17 00:00:00 2001 From: Hemang Date: Mon, 20 May 2019 11:59:18 +0530 Subject: [PATCH 1/6] Add_maximum_bytes_billed_to_magics --- bigquery/google/cloud/bigquery/magics.py | 28 ++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 6bd1c45dfcd5..fa3d2b92b2f9 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -161,6 +161,7 @@ def __init__(self): self._project = None self._connection = None self._use_bqstorage_api = None + self._maximum_bytes_billed = None @property def credentials(self): @@ -237,6 +238,32 @@ def use_bqstorage_api(self): def use_bqstorage_api(self, value): self._use_bqstorage_api = value + @property + def maximum_bytes_billed(self): + """int: Maximum bytes to be billed for this job or :data:`None` if not set. + + Example: + Manually setting the context maximum_bytes_billed: + + >>> from google.cloud.bigquery import magics + >>> magics.context.maximum_bytes_billed = '123' + + Raises: + ValueError: If the parameters are invalid. + """ + return self._maximum_bytes_billed + + @maximum_bytes_billed.setter + def maximum_bytes_billed(self, value): + try: + value = int(value) + self._maximum_bytes_billed = value + job_config = bigquery.job.QueryJobConfig() + job_config.maximum_bytes_billed = self._maximum_bytes_billed + + except: + raise ValueError("value is not a valid integer.") + context = Context() @@ -372,6 +399,7 @@ def _cell_magic(line, query): job_config = bigquery.job.QueryJobConfig() job_config.query_parameters = params job_config.use_legacy_sql = args.use_legacy_sql + job_config.maximum_bytes_billed = context.maximum_bytes_billed query_job = _run_query(client, query, job_config) if not args.verbose: From 629d34b0cc47f0da147d92d51d4a310b22c464c2 Mon Sep 17 00:00:00 2001 From: Hemang Date: Tue, 21 May 2019 13:06:07 +0530 Subject: [PATCH 2/6] Add_maximum_bytes_billed_to_magics and test cases --- bigquery/google/cloud/bigquery/magics.py | 23 ++++- bigquery/tests/unit/test_magics.py | 102 +++++++++++++++++++++++ 2 files changed, 122 insertions(+), 3 deletions(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index fa3d2b92b2f9..25c706275f9a 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -161,7 +161,7 @@ def __init__(self): self._project = None self._connection = None self._use_bqstorage_api = None - self._maximum_bytes_billed = None + self._maximum_bytes_billed = 0 @property def credentials(self): @@ -261,7 +261,7 @@ def maximum_bytes_billed(self, value): job_config = bigquery.job.QueryJobConfig() job_config.maximum_bytes_billed = self._maximum_bytes_billed - except: + except Exception: raise ValueError("value is not a valid integer.") @@ -318,6 +318,13 @@ def _run_query(client, query, job_config=None): default=None, help=("Project to use for executing this query. Defaults to the context project."), ) +@magic_arguments.argument( + "--maximum_bytes_billed", + default=None, + help=( + "maximum_bytes_billed to use for executing this query. Defaults to the context maximum_bytes_billed." + ), +) @magic_arguments.argument( "--use_legacy_sql", action="store_true", @@ -399,7 +406,17 @@ def _cell_magic(line, query): job_config = bigquery.job.QueryJobConfig() job_config.query_parameters = params job_config.use_legacy_sql = args.use_legacy_sql - job_config.maximum_bytes_billed = context.maximum_bytes_billed + + if args.maximum_bytes_billed == "None": + job_config.maximum_bytes_billed = 0 + elif args.maximum_bytes_billed is not None: + try: + value = int(args.maximum_bytes_billed) + job_config.maximum_bytes_billed = value + except Exception: + raise ValueError("value is not a valid integer.") + else: + job_config.maximum_bytes_billed = context.maximum_bytes_billed query_job = _run_query(client, query, job_config) if not args.verbose: diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 70848cbcae64..98899ddbe2c3 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -548,6 +548,88 @@ def test_bigquery_magic_without_bqstorage(monkeypatch): assert isinstance(return_value, pandas.DataFrame) +@pytest.mark.usefixtures("ipython_interactive") +def test_maximum_bytes_billed_w_int_magic(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._project = None + + bqstorage_mock = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient + ) + + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + default_patch = mock.patch( + "google.auth.default", return_value=(credentials_mock, "general-project") + ) + run_query_patch = mock.patch( + "google.cloud.bigquery.magics._run_query", autospec=True + ) + + sql = "SELECT 17 AS num" + result = pandas.DataFrame([17], columns=["num"]) + query_job_mock = mock.create_autospec( + google.cloud.bigquery.job.QueryJob, instance=True + ) + query_job_mock.to_dataframe.return_value = result + with run_query_patch as run_query_mock, default_patch: + run_query_mock.return_value = query_job_mock + return_value = ip.run_cell_magic("bigquery", "--maximum_bytes_billed=123456789", sql) + + bqstorage_mock.assert_not_called() + query_job_mock.to_dataframe.assert_called_once_with(bqstorage_client=None) + assert isinstance(return_value, pandas.DataFrame) + + +@pytest.mark.usefixtures("ipython_interactive") +def test_maximum_bytes_billed_w_string_params(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._project = None + + sql = "SELECT 17 AS num" + + with pytest.raises(ValueError): + ip.run_cell_magic("bigquery", "--maximum_bytes_billed=abc", sql) + + +@pytest.mark.usefixtures("ipython_interactive") +def test_maximum_bytes_billed_w_none__magic(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._project = None + + bqstorage_mock = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient + ) + + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + default_patch = mock.patch( + "google.auth.default", return_value=(credentials_mock, "general-project") + ) + run_query_patch = mock.patch( + "google.cloud.bigquery.magics._run_query", autospec=True + ) + + sql = "SELECT 17 AS num" + result = pandas.DataFrame([17], columns=["num"]) + query_job_mock = mock.create_autospec( + google.cloud.bigquery.job.QueryJob, instance=True + ) + query_job_mock.to_dataframe.return_value = result + with run_query_patch as run_query_mock, default_patch: + run_query_mock.return_value = query_job_mock + return_value = ip.run_cell_magic("bigquery", "--maximum_bytes_billed=None", sql) + + bqstorage_mock.assert_not_called() + query_job_mock.to_dataframe.assert_called_once_with(bqstorage_client=None) + assert isinstance(return_value, pandas.DataFrame) + + @pytest.mark.usefixtures("ipython_interactive") def test_bigquery_magic_with_project(): ip = IPython.get_ipython() @@ -652,3 +734,23 @@ def test_bigquery_magic_with_improperly_formatted_params(): with pytest.raises(SyntaxError): ip.run_cell_magic("bigquery", "--params {17}", sql) + + +def test_maximum_bytes_billed_set_value(): + """When Application Default Credentials are set, the context credentials + will be created the first time it is called + """ + + from google.cloud.bigquery import QueryJobConfig + job_config = QueryJobConfig() + magics.context.maximum_bytes_billed = 1234567489 + assert job_config.maximum_bytes_billed == magics.context.maximum_bytes_billed + + +def test_maximum_bytes_billed_set_string(): + """When Application Default Credentials are set, the context credentials + will be created the first time it is called + """ + with pytest.raises(ValueError): + magics.context.maximum_bytes_billed = "abc" + From 30f8326092e5bc16d1c4f3e9b2b1f36beedec9f8 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 31 May 2019 16:55:31 -0700 Subject: [PATCH 3/6] Add default_query_job_config property. Test API params get set. --- bigquery/google/cloud/bigquery/magics.py | 52 ++++--- bigquery/tests/unit/test_magics.py | 166 +++++++++++++---------- 2 files changed, 120 insertions(+), 98 deletions(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 25c706275f9a..0acde4f21b5f 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -161,7 +161,7 @@ def __init__(self): self._project = None self._connection = None self._use_bqstorage_api = None - self._maximum_bytes_billed = 0 + self._default_query_job_config = bigquery.QueryJobConfig() @property def credentials(self): @@ -239,30 +239,26 @@ def use_bqstorage_api(self, value): self._use_bqstorage_api = value @property - def maximum_bytes_billed(self): - """int: Maximum bytes to be billed for this job or :data:`None` if not set. + def default_query_job_config(self): + """google.cloud.bigquery.job.QueryJobConfig: Default job + configuration for queries. + + The context's :class:`~google.cloud.bigquery.job.QueryJobConfig` is + used for queries. Some properties can be overridden with arguments to + the magics. Example: - Manually setting the context maximum_bytes_billed: + Manually setting the default value for ``maximum_bytes_billed`` + to 100 MB: >>> from google.cloud.bigquery import magics - >>> magics.context.maximum_bytes_billed = '123' - - Raises: - ValueError: If the parameters are invalid. + >>> magics.context.default_query_job_config.maximum_bytes_billed = 100000000 """ - return self._maximum_bytes_billed + return self._default_query_job_config - @maximum_bytes_billed.setter - def maximum_bytes_billed(self, value): - try: - value = int(value) - self._maximum_bytes_billed = value - job_config = bigquery.job.QueryJobConfig() - job_config.maximum_bytes_billed = self._maximum_bytes_billed - - except Exception: - raise ValueError("value is not a valid integer.") + @default_query_job_config.setter + def default_query_job_config(self, value): + self._default_query_job_config = value context = Context() @@ -322,7 +318,8 @@ def _run_query(client, query, job_config=None): "--maximum_bytes_billed", default=None, help=( - "maximum_bytes_billed to use for executing this query. Defaults to the context maximum_bytes_billed." + "maximum_bytes_billed to use for executing this query. Defaults to " + "the context default_query_job_config.maximum_bytes_billed." ), ) @magic_arguments.argument( @@ -397,7 +394,11 @@ def _cell_magic(line, query): ) project = args.project or context.project - client = bigquery.Client(project=project, credentials=context.credentials) + client = bigquery.Client( + project=project, + credentials=context.credentials, + default_query_job_config=context.default_query_job_config, + ) if context._connection: client._connection = context._connection bqstorage_client = _make_bqstorage_client( @@ -410,13 +411,8 @@ def _cell_magic(line, query): if args.maximum_bytes_billed == "None": job_config.maximum_bytes_billed = 0 elif args.maximum_bytes_billed is not None: - try: - value = int(args.maximum_bytes_billed) - job_config.maximum_bytes_billed = value - except Exception: - raise ValueError("value is not a valid integer.") - else: - job_config.maximum_bytes_billed = context.maximum_bytes_billed + value = int(args.maximum_bytes_billed) + job_config.maximum_bytes_billed = value query_job = _run_query(client, query, job_config) if not args.verbose: diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 98899ddbe2c3..f685b5882053 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import re -import mock -import six from concurrent import futures +import mock import pytest +import six try: import pandas @@ -63,6 +64,26 @@ def ipython_interactive(request, ipython): yield ipython +JOB_REFERENCE_RESOURCE = {"projectId": "its-a-project-eh", "jobId": "some-random-id"} +TABLE_REFERENCE_RESOURCE = { + "projectId": "its-a-project-eh", + "datasetId": "ds", + "tableId": "persons", +} +QUERY_RESOURCE = { + "jobReference": JOB_REFERENCE_RESOURCE, + "configuration": { + "query": { + "destinationTable": TABLE_REFERENCE_RESOURCE, + "query": "SELECT 42 FROM `life.the_universe.and_everything`;", + "queryParameters": [], + "useLegacySql": False, + } + }, + "status": {"state": "DONE"}, +} + + def test_context_credentials_auto_set_w_application_default_credentials(): """When Application Default Credentials are set, the context credentials will be created the first time it is called @@ -117,22 +138,13 @@ def test_context_connection_can_be_overriden(): 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" - 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"}, - } + resource = copy.deepcopy(QUERY_RESOURCE) + resource["jobReference"] = job_reference + resource["configuration"]["query"]["query"] = query data = {"jobReference": job_reference, "totalRows": 0, "rows": []} conn = magics.context._connection = make_connection(resource, data) @@ -170,22 +182,13 @@ def test_context_no_connection(): 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" - 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"}, - } + 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) @@ -549,7 +552,7 @@ def test_bigquery_magic_without_bqstorage(monkeypatch): @pytest.mark.usefixtures("ipython_interactive") -def test_maximum_bytes_billed_w_int_magic(): +def test_bigquery_magic_w_maximum_bytes_billed(): ip = IPython.get_ipython() ip.extension_manager.load_extension("google.cloud.bigquery") magics.context._project = None @@ -576,7 +579,9 @@ def test_maximum_bytes_billed_w_int_magic(): query_job_mock.to_dataframe.return_value = result with run_query_patch as run_query_mock, default_patch: run_query_mock.return_value = query_job_mock - return_value = ip.run_cell_magic("bigquery", "--maximum_bytes_billed=123456789", sql) + return_value = ip.run_cell_magic( + "bigquery", "--maximum_bytes_billed=123456789", sql + ) bqstorage_mock.assert_not_called() query_job_mock.to_dataframe.assert_called_once_with(bqstorage_client=None) @@ -584,7 +589,7 @@ def test_maximum_bytes_billed_w_int_magic(): @pytest.mark.usefixtures("ipython_interactive") -def test_maximum_bytes_billed_w_string_params(): +def test_bigquery_magic_w_maximum_bytes_billed_invalid(): ip = IPython.get_ipython() ip.extension_manager.load_extension("google.cloud.bigquery") magics.context._project = None @@ -595,39 +600,80 @@ def test_maximum_bytes_billed_w_string_params(): ip.run_cell_magic("bigquery", "--maximum_bytes_billed=abc", sql) +@pytest.mark.parametrize( + "param_value,expected", [("987654321", "987654321"), ("None", "0")] +) @pytest.mark.usefixtures("ipython_interactive") -def test_maximum_bytes_billed_w_none__magic(): +def test_bigquery_magic_w_maximum_bytes_billed_overrides_context(param_value, expected): ip = IPython.get_ipython() ip.extension_manager.load_extension("google.cloud.bigquery") magics.context._project = None - bqstorage_mock = mock.create_autospec( - bigquery_storage_v1beta1.BigQueryStorageClient - ) + # Set the default maximum bytes billed, so we know it's overridable by the param. + magics.context.default_query_job_config.maximum_bytes_billed = 1234567 + 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") ) - run_query_patch = mock.patch( - "google.cloud.bigquery.magics._run_query", autospec=True + 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", "--maximum_bytes_billed={}".format(param_value), query + ) - sql = "SELECT 17 AS num" - result = pandas.DataFrame([17], columns=["num"]) - query_job_mock = mock.create_autospec( - google.cloud.bigquery.job.QueryJob, instance=True + _, req = conn.api_request.call_args_list[0] + sent_config = req["data"]["configuration"]["query"] + assert sent_config["maximumBytesBilled"] == expected + + +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_w_maximum_bytes_billed_w_context(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._project = None + + magics.context.default_query_job_config.maximum_bytes_billed = 1234567 + + 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 ) - query_job_mock.to_dataframe.return_value = result - with run_query_patch as run_query_mock, default_patch: - run_query_mock.return_value = query_job_mock - return_value = ip.run_cell_magic("bigquery", "--maximum_bytes_billed=None", sql) + 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) - bqstorage_mock.assert_not_called() - query_job_mock.to_dataframe.assert_called_once_with(bqstorage_client=None) - assert isinstance(return_value, pandas.DataFrame) + _, req = conn.api_request.call_args_list[0] + sent_config = req["data"]["configuration"]["query"] + assert sent_config["maximumBytesBilled"] == "1234567" @pytest.mark.usefixtures("ipython_interactive") @@ -734,23 +780,3 @@ def test_bigquery_magic_with_improperly_formatted_params(): with pytest.raises(SyntaxError): ip.run_cell_magic("bigquery", "--params {17}", sql) - - -def test_maximum_bytes_billed_set_value(): - """When Application Default Credentials are set, the context credentials - will be created the first time it is called - """ - - from google.cloud.bigquery import QueryJobConfig - job_config = QueryJobConfig() - magics.context.maximum_bytes_billed = 1234567489 - assert job_config.maximum_bytes_billed == magics.context.maximum_bytes_billed - - -def test_maximum_bytes_billed_set_string(): - """When Application Default Credentials are set, the context credentials - will be created the first time it is called - """ - with pytest.raises(ValueError): - magics.context.maximum_bytes_billed = "abc" - From a4de3335db08eba44ba4a704032e0a4539f2afc2 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 31 May 2019 17:01:55 -0700 Subject: [PATCH 4/6] Add test for default_query_job_config setter --- bigquery/tests/unit/test_magics.py | 44 ++++++++++++++++++++++++++++-- 1 file changed, 41 insertions(+), 3 deletions(-) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index f685b5882053..19adae1be456 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -38,6 +38,7 @@ from google.cloud import bigquery_storage_v1beta1 except ImportError: # pragma: NO COVER bigquery_storage_v1beta1 = None +from google.cloud.bigquery import job from google.cloud.bigquery import table from google.cloud.bigquery import magics from tests.unit.helpers import make_connection @@ -642,12 +643,49 @@ def test_bigquery_magic_w_maximum_bytes_billed_overrides_context(param_value, ex @pytest.mark.usefixtures("ipython_interactive") -def test_bigquery_magic_w_maximum_bytes_billed_w_context(): +def test_bigquery_magic_w_maximum_bytes_billed_w_context_inplace(): ip = IPython.get_ipython() ip.extension_manager.load_extension("google.cloud.bigquery") magics.context._project = None - magics.context.default_query_job_config.maximum_bytes_billed = 1234567 + magics.context.default_query_job_config.maximum_bytes_billed = 1337 + + 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"] == "1337" + + +@pytest.mark.usefixtures("ipython_interactive") +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) @@ -673,7 +711,7 @@ def test_bigquery_magic_w_maximum_bytes_billed_w_context(): _, req = conn.api_request.call_args_list[0] sent_config = req["data"]["configuration"]["query"] - assert sent_config["maximumBytesBilled"] == "1234567" + assert sent_config["maximumBytesBilled"] == "10203" @pytest.mark.usefixtures("ipython_interactive") From 266de84a5efc2fba03fa2f1fc29d9bdf78985c16 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Mon, 3 Jun 2019 10:37:31 -0700 Subject: [PATCH 5/6] Remove redundant maximum_bytes_billed test. --- bigquery/tests/unit/test_magics.py | 37 ------------------------------ 1 file changed, 37 deletions(-) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 19adae1be456..025f0f86a077 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -552,43 +552,6 @@ def test_bigquery_magic_without_bqstorage(monkeypatch): assert isinstance(return_value, pandas.DataFrame) -@pytest.mark.usefixtures("ipython_interactive") -def test_bigquery_magic_w_maximum_bytes_billed(): - ip = IPython.get_ipython() - ip.extension_manager.load_extension("google.cloud.bigquery") - magics.context._project = None - - bqstorage_mock = mock.create_autospec( - bigquery_storage_v1beta1.BigQueryStorageClient - ) - - credentials_mock = mock.create_autospec( - google.auth.credentials.Credentials, instance=True - ) - default_patch = mock.patch( - "google.auth.default", return_value=(credentials_mock, "general-project") - ) - run_query_patch = mock.patch( - "google.cloud.bigquery.magics._run_query", autospec=True - ) - - sql = "SELECT 17 AS num" - result = pandas.DataFrame([17], columns=["num"]) - query_job_mock = mock.create_autospec( - google.cloud.bigquery.job.QueryJob, instance=True - ) - query_job_mock.to_dataframe.return_value = result - with run_query_patch as run_query_mock, default_patch: - run_query_mock.return_value = query_job_mock - return_value = ip.run_cell_magic( - "bigquery", "--maximum_bytes_billed=123456789", sql - ) - - bqstorage_mock.assert_not_called() - query_job_mock.to_dataframe.assert_called_once_with(bqstorage_client=None) - assert isinstance(return_value, pandas.DataFrame) - - @pytest.mark.usefixtures("ipython_interactive") def test_bigquery_magic_w_maximum_bytes_billed_invalid(): ip = IPython.get_ipython() From 7ebfcc581b419857178502c231363ec420cafc7b Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Mon, 3 Jun 2019 11:23:39 -0700 Subject: [PATCH 6/6] Use for loop instead of `all` and generator expression. Maybe that'll fix the coverage issues? --- bigquery/tests/unit/test_magics.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 025f0f86a077..f3e64a46faca 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -243,7 +243,8 @@ def test__run_query(): assert updates[0] == expected_first_line execution_updates = updates[1:-1] assert len(execution_updates) == 3 # one update per API response - assert all(re.match("Query executing: .*s", line) for line in execution_updates) + for line in execution_updates: + assert re.match("Query executing: .*s", line) assert re.match("Query complete after .*s", updates[-1])