This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new dac3da4ccfd Add missing tests for google.cloud.utils BigQuery helpers
(#72276) (#72448)
dac3da4ccfd is described below
commit dac3da4ccfd68afb3feb3683c3d2ea420b60eb0a
Author: arunjvattoly <[email protected]>
AuthorDate: Thu Oct 29 02:55:51 2026 -0500
Add missing tests for google.cloud.utils BigQuery helpers (#72276) (#72448)
Co-authored-by: Shahar Epstein <[email protected]>
---
.../tests/unit/always/test_project_structure.py | 2 -
.../tests/unit/google/cloud/utils/test_bigquery.py | 58 ++++++++++++++++
.../google/cloud/utils/test_bigquery_get_data.py | 77 ++++++++++++++++++++++
3 files changed, 135 insertions(+), 2 deletions(-)
diff --git a/airflow-core/tests/unit/always/test_project_structure.py
b/airflow-core/tests/unit/always/test_project_structure.py
index aa43f5013f1..cbade3332f6 100644
--- a/airflow-core/tests/unit/always/test_project_structure.py
+++ b/airflow-core/tests/unit/always/test_project_structure.py
@@ -101,8 +101,6 @@ class TestProjectStructure:
"providers/google/tests/unit/google/cloud/sensors/vertex_ai/test_feature_store.py",
"providers/google/tests/unit/google/cloud/transfers/test_bigquery_to_sql.py",
"providers/google/tests/unit/google/cloud/transfers/test_presto_to_gcs.py",
- "providers/google/tests/unit/google/cloud/utils/test_bigquery.py",
-
"providers/google/tests/unit/google/cloud/utils/test_bigquery_get_data.py",
"providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_adls.py",
"providers/snowflake/tests/unit/snowflake/triggers/test_snowflake_trigger.py",
]
diff --git a/providers/google/tests/unit/google/cloud/utils/test_bigquery.py
b/providers/google/tests/unit/google/cloud/utils/test_bigquery.py
new file mode 100644
index 00000000000..786bda7c76d
--- /dev/null
+++ b/providers/google/tests/unit/google/cloud/utils/test_bigquery.py
@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import pytest
+
+from airflow.providers.google.cloud.utils.bigquery import bq_cast,
convert_job_id
+
+
+class TestBqCast:
+ @pytest.mark.parametrize(
+ ("value", "bq_type", "expected"),
+ [
+ (None, "INTEGER", None),
+ ("123", "INTEGER", 123),
+ ("123.45", "FLOAT", 123.45),
+ ("123.45", "TIMESTAMP", 123.45),
+ ("true", "BOOLEAN", True),
+ ("false", "BOOLEAN", False),
+ ("raw_string", "UNKNOWN_TYPE", "raw_string"),
+ ],
+ )
+ def test_bq_cast(self, value, bq_type, expected):
+ result = bq_cast(value, bq_type)
+ assert result == expected
+ assert type(result) is type(expected)
+
+ def test_bq_cast_boolean_invalid_raises_value_error(self):
+ with pytest.raises(ValueError, match="invalid must have value 'true'
or 'false'"):
+ bq_cast("invalid", "BOOLEAN")
+
+
+class TestConvertJobId:
+ @pytest.mark.parametrize(
+ ("job_id", "location", "expected"),
+ [
+ ("job123", "EU", "proj:EU:job123"),
+ ("job123", None, "proj:US:job123"),
+ (["job1", "job2"], "EU", ["proj:EU:job1", "proj:EU:job2"]),
+ ],
+ )
+ def test_convert_job_id(self, job_id, location, expected):
+ assert convert_job_id(job_id=job_id, project_id="proj",
location=location) == expected
diff --git
a/providers/google/tests/unit/google/cloud/utils/test_bigquery_get_data.py
b/providers/google/tests/unit/google/cloud/utils/test_bigquery_get_data.py
new file mode 100644
index 00000000000..fd907bf26db
--- /dev/null
+++ b/providers/google/tests/unit/google/cloud/utils/test_bigquery_get_data.py
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import logging
+from unittest import mock
+
+import pytest
+from google.cloud.bigquery.table import Row, RowIterator
+
+from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
+from airflow.providers.google.cloud.utils.bigquery_get_data import
bigquery_get_data
+
+
+class TestBigqueryGetData:
+ def test_yields_successive_batches(self):
+ mock_hook = mock.MagicMock(spec=BigQueryHook)
+ mock_logger = mock.MagicMock(spec=logging.Logger)
+ rows = [mock.MagicMock(spec=Row) for _ in range(3)]
+ for row, value in zip(rows, "ABC"):
+ row.values.return_value = (value,)
+ mock_hook.list_rows.side_effect = [rows[:2], rows[2:], []]
+
+ results = list(
+ bigquery_get_data(
+ logger=mock_logger,
+ dataset_id="my_dataset",
+ table_id="my_table",
+ big_query_hook=mock_hook,
+ batch_size=2,
+ selected_fields=["col1"],
+ )
+ )
+
+ assert results == [[("A",), ("B",)], [("C",)]]
+ assert mock_hook.list_rows.mock_calls == [
+ mock.call(
+ dataset_id="my_dataset",
+ table_id="my_table",
+ max_results=2,
+ selected_fields=["col1"],
+ start_index=start_index,
+ )
+ for start_index in (0, 2, 4)
+ ]
+ mock_logger.info.assert_any_call("Job Finished")
+
+ def test_raises_type_error_for_row_iterator(self):
+ mock_hook = mock.MagicMock(spec=BigQueryHook)
+ mock_hook.list_rows.return_value = mock.MagicMock(spec=RowIterator)
+
+ with pytest.raises(TypeError, match="returns iterator when
return_iterator=False"):
+ list(
+ bigquery_get_data(
+ logger=mock.MagicMock(spec=logging.Logger),
+ dataset_id="my_dataset",
+ table_id="my_table",
+ big_query_hook=mock_hook,
+ batch_size=1,
+ selected_fields=["col1"],
+ )
+ )
