added tests for --max_results magic option

bigquery/google/cloud/bigquery/magics.py

@@ -320,7 +320,6 @@ def _run_query(client, query, job_config=None, max_results=None):
     default=None,
     help=("Project to use for executing this query. Defaults to the context project."),
 )
-
 @magic_arguments.argument(
     "--max_results",
     default=None,
@@ -329,7 +328,6 @@ def _run_query(client, query, job_config=None, max_results=None):
         "Defaults to returning all rows."
     ),
 )
-
 @magic_arguments.argument(
     "--maximum_bytes_billed",
     default=None,
@@ -449,7 +447,9 @@ def _cell_magic(line, query):
 
     error = None
     try:
-        query_job = _run_query(client, query, job_config=job_config, max_results=max_results)
+        query_job = _run_query(
+            client, query, job_config=job_config, max_results=max_results
+        )
     except Exception as ex:
         error = str(ex)
 

bigquery/tests/unit/test_magics.py

@@ -414,7 +414,7 @@ def test_bigquery_magic_with_legacy_sql():
     with run_query_patch as run_query_mock:
         ip.run_cell_magic("bigquery", "--use_legacy_sql", "SELECT 17 AS num")
 
-        job_config_used = run_query_mock.call_args_list[0][1]['job_config']
+        job_config_used = run_query_mock.call_args_list[0][1]["job_config"]
         assert job_config_used.use_legacy_sql is True
 
 
@@ -645,6 +645,57 @@ def test_bigquery_magic_without_bqstorage(monkeypatch):
     assert isinstance(return_value, pandas.DataFrame)
 
 
+@pytest.mark.usefixtures("ipython_interactive")
+def test_bigquery_magic_w_max_results_invalid():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context._project = None
+
+    credentials_mock = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+    default_patch = mock.patch(
+        "google.auth.default", return_value=(credentials_mock, "general-project")
+    )
+    client_query_patch = mock.patch(
+        "google.cloud.bigquery.client.Client.query", autospec=True
+    )
+
+    sql = "SELECT 17 AS num"
+
+    with pytest.raises(ValueError), default_patch, client_query_patch:
+        ip.run_cell_magic("bigquery", "--max_results=abc", sql)
+
+
+@pytest.mark.usefixtures("ipython_interactive")
+def test_bigquery_magic_w_max_results_valid_calls_queryjob_result():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context._project = None
+
+    credentials_mock = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+    default_patch = mock.patch(
+        "google.auth.default", return_value=(credentials_mock, "general-project")
+    )
+    client_query_patch = mock.patch(
+        "google.cloud.bigquery.client.Client.query", autospec=True
+    )
+
+    sql = "SELECT 17 AS num"
+
+    query_job_mock = mock.create_autospec(
+        google.cloud.bigquery.job.QueryJob, instance=True
+    )
+
+    with client_query_patch as client_query_mock, default_patch:
+        client_query_mock.return_value = query_job_mock
+        ip.run_cell_magic("bigquery", "--max_results=5", sql)
+
+        query_job_mock.result.assert_called_once_with(timeout=0.5, max_results=5)
+
+
 @pytest.mark.usefixtures("ipython_interactive")
 def test_bigquery_magic_dryrun_option_sets_job_config():
     ip = IPython.get_ipython()
@@ -662,7 +713,7 @@ def test_bigquery_magic_dryrun_option_sets_job_config():
     with run_query_patch as run_query_mock:
         ip.run_cell_magic("bigquery", "--dry_run", sql)
 
-        job_config_used = run_query_mock.call_args_list[0][1]['job_config']
+        job_config_used = run_query_mock.call_args_list[0][1]["job_config"]
         assert job_config_used.dry_run is True
 
 
@@ -924,7 +975,9 @@ def test_bigquery_magic_with_string_params():
         run_query_mock.return_value = query_job_mock
 
         ip.run_cell_magic("bigquery", 'params_string_df --params {"num":17}', sql)
-        run_query_mock.assert_called_once_with(mock.ANY, sql.format(num=17), mock.ANY, max_results=None)
+        run_query_mock.assert_called_once_with(
+            mock.ANY, sql.format(num=17), mock.ANY, max_results=None
+        )
 
     assert "params_string_df" in ip.user_ns  # verify that the variable exists
     df = ip.user_ns["params_string_df"]
@@ -959,7 +1012,9 @@ def test_bigquery_magic_with_dict_params():
         # Insert dictionary into user namespace so that it can be expanded
         ip.user_ns["params"] = params
         ip.run_cell_magic("bigquery", "params_dict_df --params $params", sql)
-        run_query_mock.assert_called_once_with(mock.ANY, sql.format(num=17), mock.ANY, max_results=None)
+        run_query_mock.assert_called_once_with(
+            mock.ANY, sql.format(num=17), mock.ANY, max_results=None
+        )
 
     assert "params_dict_df" in ip.user_ns  # verify that the variable exists
     df = ip.user_ns["params_dict_df"]