|
21 | 21 | import tempfile
|
22 | 22 | import textwrap
|
23 | 23 |
|
24 |
| -from google.api_core.exceptions import BadRequest, InvalidArgument, NotFound |
| 24 | +import google.api_core.exceptions |
25 | 25 | from google.cloud import bigquery, storage
|
26 | 26 | import pandas
|
27 | 27 | import pytest
|
@@ -849,7 +849,7 @@ def square(x):
|
849 | 849 | expected_remote_function = f"{dataset_id}.{rf_name}"
|
850 | 850 |
|
851 | 851 | # Initially the expected BQ remote function should not exist
|
852 |
| - with pytest.raises(NotFound): |
| 852 | + with pytest.raises(google.api_core.exceptions.NotFound): |
853 | 853 | session.bqclient.get_routine(expected_remote_function)
|
854 | 854 |
|
855 | 855 | # Create the remote function with the name provided explicitly
|
@@ -981,7 +981,7 @@ def test_internal(rf, udf):
|
981 | 981 | expected_remote_function = f"{dataset_id}.{rf_name}"
|
982 | 982 |
|
983 | 983 | # Initially the expected BQ remote function should not exist
|
984 |
| - with pytest.raises(NotFound): |
| 984 | + with pytest.raises(google.api_core.exceptions.NotFound): |
985 | 985 | session.bqclient.get_routine(expected_remote_function)
|
986 | 986 |
|
987 | 987 | # Create a new remote function with the name provided explicitly
|
@@ -1198,7 +1198,8 @@ def square(x):
|
1198 | 1198 | scalars_df, _ = scalars_dfs
|
1199 | 1199 |
|
1200 | 1200 | with pytest.raises(
|
1201 |
| - BadRequest, match="400.*errorMessage.*unsupported operand type" |
| 1201 | + google.api_core.exceptions.BadRequest, |
| 1202 | + match="400.*errorMessage.*unsupported operand type", |
1202 | 1203 | ):
|
1203 | 1204 | # int64_col has nulls which should cause error in square
|
1204 | 1205 | scalars_df["int64_col"].apply(square).to_pandas()
|
@@ -1431,20 +1432,6 @@ def square_num(x):
|
1431 | 1432 | )
|
1432 | 1433 |
|
1433 | 1434 |
|
1434 |
| -def test_remote_function_via_session_vpc_invalid(session): |
1435 |
| - with pytest.raises( |
1436 |
| - InvalidArgument, match="400.*Serverless VPC Access connector is not found" |
1437 |
| - ): |
1438 |
| - |
1439 |
| - @session.remote_function( |
1440 |
| - [int], int, reuse=False, cloud_function_vpc_connector="does-not-exist" |
1441 |
| - ) |
1442 |
| - def square_num(x): |
1443 |
| - if x is None: |
1444 |
| - return x |
1445 |
| - return x * x |
1446 |
| - |
1447 |
| - |
1448 | 1435 | @pytest.mark.parametrize(
|
1449 | 1436 | ("max_batching_rows"),
|
1450 | 1437 | [
|
|
0 commit comments