Skip to content

Commit

Permalink
Improve cell magic parser test coverage
Browse files Browse the repository at this point in the history
  • Loading branch information
plamut committed Aug 27, 2020
1 parent d95802b commit b514850
Showing 1 changed file with 229 additions and 7 deletions.
236 changes: 229 additions & 7 deletions tests/unit/test_magics.py
Expand Up @@ -43,7 +43,6 @@
from google.cloud import bigquery
from google.cloud.bigquery import job
from google.cloud.bigquery import table
from google.cloud.bigquery.ipython_magics import line_arg_parser as lap
from google.cloud.bigquery.ipython_magics import magics
from tests.unit.helpers import make_connection
from test_utils.imports import maybe_fail_import
Expand Down Expand Up @@ -1165,6 +1164,37 @@ def test_bigquery_magic_with_project():
assert magics.context.project == "general-project"


@pytest.mark.usefixtures("ipython_interactive")
def test_bigquery_magic_with_multiple_options():
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")
)
run_query_patch = mock.patch(
"google.cloud.bigquery.ipython_magics.magics._run_query", autospec=True
)
with run_query_patch as run_query_mock, default_patch:
ip.run_cell_magic(
"bigquery",
"--project=specific-project --use_legacy_sql --maximum_bytes_billed 1024",
"SELECT 17 as num",
)

args, kwargs = run_query_mock.call_args
client_used = args[0]
assert client_used.project == "specific-project"

job_config_used = kwargs["job_config"]
assert job_config_used.use_legacy_sql
assert job_config_used.maximum_bytes_billed == 1024


@pytest.mark.usefixtures("ipython_interactive")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
def test_bigquery_magic_with_string_params():
Expand All @@ -1176,7 +1206,9 @@ def test_bigquery_magic_with_string_params():

sql = "SELECT @num AS num"
result = pandas.DataFrame([17], columns=["num"])
assert "params_string_df" not in ip.user_ns

if "params_dict_df" in ip.user_ns:
del ip.user_ns["params_dict_df"]

run_query_patch = mock.patch(
"google.cloud.bigquery.ipython_magics.magics._run_query", autospec=True
Expand Down Expand Up @@ -1207,9 +1239,13 @@ def test_bigquery_magic_with_dict_params():
google.auth.credentials.Credentials, instance=True
)

sql = "SELECT @num AS num"
result = pandas.DataFrame([17], columns=["num"])
assert "params_dict_df" not in ip.user_ns
sql = "SELECT @num AS num, @tricky_value as tricky_value"
result = pandas.DataFrame(
[(17, '--params "value"')], columns=["num", "tricky_value"]
)

if "params_dict_df" in ip.user_ns:
del ip.user_ns["params_dict_df"]

run_query_patch = mock.patch(
"google.cloud.bigquery.ipython_magics.magics._run_query", autospec=True
Expand All @@ -1221,7 +1257,7 @@ def test_bigquery_magic_with_dict_params():
with run_query_patch as run_query_mock:
run_query_mock.return_value = query_job_mock

params = {"num": 17}
params = {"num": 17, "tricky_value": '--params "value"'}
# 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)
Expand All @@ -1233,6 +1269,192 @@ def test_bigquery_magic_with_dict_params():
assert len(df) == len(result) # verify row count
assert list(df) == list(result) # verify column names

assert df["num"][0] == 17
assert df["tricky_value"][0] == '--params "value"'


@pytest.mark.usefixtures("ipython_interactive")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
def test_bigquery_magic_with_dict_params_nonexisting():
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
magics.context.credentials = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)

sql = "SELECT @foo AS foo"

with pytest.raises(NameError, match=r".*undefined variable.*unknown_name.*"):
ip.run_cell_magic("bigquery", "params_dict_df --params $unknown_name", sql)


@pytest.mark.usefixtures("ipython_interactive")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
def test_bigquery_magic_with_dict_params_incorrect_syntax():
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
magics.context.credentials = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)

sql = "SELECT @foo AS foo"

with pytest.raises(SyntaxError, match=r".*--params.*"):
cell_magic_args = "params_dict_df --params {'foo': 1; 'bar': 2}"
ip.run_cell_magic("bigquery", cell_magic_args, sql)


@pytest.mark.usefixtures("ipython_interactive")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
def test_bigquery_magic_with_dict_params_duplicate():
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
magics.context.credentials = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)

sql = "SELECT @foo AS foo"

with pytest.raises(ValueError, match=r"Duplicate --params option\."):
cell_magic_args = (
"params_dict_df --params {'foo': 1} " "--verbose " "--params {'bar': 2} "
)
ip.run_cell_magic("bigquery", cell_magic_args, sql)


@pytest.mark.usefixtures("ipython_interactive")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
def test_bigquery_magic_with_option_value_incorrect():
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
magics.context.credentials = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)

sql = "SELECT @foo AS foo"

exc_pattern = r".*[Uu]nrecognized input.*option values correct\?.*"
with pytest.raises(ValueError, match=exc_pattern):
cell_magic_args = "params_dict_df --max_results [PLENTY!]"
ip.run_cell_magic("bigquery", cell_magic_args, sql)


@pytest.mark.usefixtures("ipython_interactive")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
def test_bigquery_magic_with_dict_params_negative_value():
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
magics.context.credentials = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)

sql = "SELECT @num AS num"
result = pandas.DataFrame([-17], columns=["num"])

if "params_dict_df" in ip.user_ns:
del ip.user_ns["params_dict_df"]

run_query_patch = mock.patch(
"google.cloud.bigquery.ipython_magics.magics._run_query", autospec=True
)
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:
run_query_mock.return_value = query_job_mock

params = {"num": -17}
# 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)

assert "params_dict_df" in ip.user_ns # verify that the variable exists
df = ip.user_ns["params_dict_df"]
assert len(df) == len(result) # verify row count
assert list(df) == list(result) # verify column names
assert df["num"][0] == -17


@pytest.mark.usefixtures("ipython_interactive")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
def test_bigquery_magic_with_dict_params_array_value():
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
magics.context.credentials = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)

sql = "SELECT @num AS num"
result = pandas.DataFrame(["foo bar", "baz quux"], columns=["array_data"])

if "params_dict_df" in ip.user_ns:
del ip.user_ns["params_dict_df"]

run_query_patch = mock.patch(
"google.cloud.bigquery.ipython_magics.magics._run_query", autospec=True
)
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:
run_query_mock.return_value = query_job_mock

params = {"array_data": ["foo bar", "baz quux"]}
# 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)

assert "params_dict_df" in ip.user_ns # verify that the variable exists
df = ip.user_ns["params_dict_df"]
assert len(df) == len(result) # verify row count
assert list(df) == list(result) # verify column names
assert list(df["array_data"]) == ["foo bar", "baz quux"]


@pytest.mark.usefixtures("ipython_interactive")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
def test_bigquery_magic_with_dict_params_tuple_value():
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
magics.context.credentials = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
)

sql = "SELECT @num AS num"
result = pandas.DataFrame(["foo bar", "baz quux"], columns=["array_data"])

if "params_dict_df" in ip.user_ns:
del ip.user_ns["params_dict_df"]

run_query_patch = mock.patch(
"google.cloud.bigquery.ipython_magics.magics._run_query", autospec=True
)
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:
run_query_mock.return_value = query_job_mock

params = {"array_data": ("foo bar", "baz quux")}
# 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)

assert "params_dict_df" in ip.user_ns # verify that the variable exists
df = ip.user_ns["params_dict_df"]
assert len(df) == len(result) # verify row count
assert list(df) == list(result) # verify column names
assert list(df["array_data"]) == ["foo bar", "baz quux"]


@pytest.mark.usefixtures("ipython_interactive")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
Expand All @@ -1245,7 +1467,7 @@ def test_bigquery_magic_with_improperly_formatted_params():

sql = "SELECT @num AS num"

with pytest.raises(lap.exceptions.QueryParamsParseError):
with pytest.raises(SyntaxError):
ip.run_cell_magic("bigquery", "--params {17}", sql)


Expand Down

0 comments on commit b514850

Please sign in to comment.