This is an automated email from the ASF dual-hosted git repository.
potiuk 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 a7342c7210e Add tests for DataFusion base classes (#71983)
a7342c7210e is described below
commit a7342c7210e502a6cb88b2c7cbc86b7a04847d2f
Author: Xinyao Zhang <[email protected]>
AuthorDate: Wed Sep 9 09:05:04 2026 -0400
Add tests for DataFusion base classes (#71983)
Cover object-store bucket parsing, format-handler configuration, and the
DataFusion exception hierarchy while removing their project-structure
exemptions.
---
.../tests/unit/always/test_project_structure.py | 3 +-
.../tests/unit/common/sql/datafusion/test_base.py | 74 ++++++++++++++++++++++
.../unit/common/sql/datafusion/test_exceptions.py | 40 ++++++++++++
3 files changed, 115 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 fe0d230deac..28eac70d585 100644
--- a/airflow-core/tests/unit/always/test_project_structure.py
+++ b/airflow-core/tests/unit/always/test_project_structure.py
@@ -97,8 +97,7 @@ class TestProjectStructure:
"providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_kubernetes_pod.py",
"providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_delete_from.py",
"providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_k8s_hashlib_wrapper.py",
-
"providers/common/sql/tests/unit/common/sql/datafusion/test_base.py",
-
"providers/common/sql/tests/unit/common/sql/datafusion/test_exceptions.py",
+ "providers/common/ai/tests/unit/common/ai/test_exceptions.py",
"providers/common/compat/tests/unit/common/compat/lineage/test_entities.py",
"providers/common/compat/tests/unit/common/compat/standard/test_operators.py",
"providers/common/compat/tests/unit/common/compat/standard/test_triggers.py",
diff --git a/providers/common/sql/tests/unit/common/sql/datafusion/test_base.py
b/providers/common/sql/tests/unit/common/sql/datafusion/test_base.py
new file mode 100644
index 00000000000..b95fac23665
--- /dev/null
+++ b/providers/common/sql/tests/unit/common/sql/datafusion/test_base.py
@@ -0,0 +1,74 @@
+# 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
+
+from typing import Any
+
+import pytest
+
+from airflow.providers.common.sql.config import DataSourceConfig, FormatType,
StorageType
+from airflow.providers.common.sql.datafusion.base import FormatHandler,
ObjectStorageProvider
+
+
+class S3ObjectStorageProvider(ObjectStorageProvider):
+ @property
+ def get_storage_type(self) -> StorageType:
+ return StorageType.S3
+
+ def create_object_store(self, path: str, connection_config=None) -> Any:
+ return None
+
+ def get_scheme(self) -> str:
+ return "s3://"
+
+
+class ParquetFormatHandler(FormatHandler):
+ @property
+ def get_format(self) -> FormatType:
+ return FormatType.PARQUET
+
+ def register_data_source_format(self, ctx) -> None:
+ pass
+
+
[email protected](
+ ("path", "expected_bucket"),
+ [
+ ("s3://example-bucket/path/to/data.parquet", "example-bucket"),
+ ("s3://example-bucket", "example-bucket"),
+ ("s3://example-bucket/", "example-bucket"),
+ ("file://example-bucket/data.parquet", None),
+ ("example-bucket/data.parquet", None),
+ ("", None),
+ ],
+)
+def test_get_bucket(path, expected_bucket):
+ assert S3ObjectStorageProvider().get_bucket(path) == expected_bucket
+
+
+def test_format_handler_stores_datasource_config():
+ datasource_config = DataSourceConfig(
+ conn_id="aws_default",
+ table_name="events",
+ uri="s3://example-bucket/events.parquet",
+ format=FormatType.PARQUET,
+ )
+
+ handler = ParquetFormatHandler(datasource_config)
+
+ assert handler.datasource_config is datasource_config
+ assert handler.get_format is FormatType.PARQUET
diff --git
a/providers/common/sql/tests/unit/common/sql/datafusion/test_exceptions.py
b/providers/common/sql/tests/unit/common/sql/datafusion/test_exceptions.py
new file mode 100644
index 00000000000..448abb1ffdc
--- /dev/null
+++ b/providers/common/sql/tests/unit/common/sql/datafusion/test_exceptions.py
@@ -0,0 +1,40 @@
+# 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.common.compat.sdk import AirflowException
+from airflow.providers.common.sql.datafusion.exceptions import (
+ FileFormatRegistrationException,
+ IcebergRegistrationException,
+ ObjectStoreCreationException,
+ QueryExecutionException,
+)
+
+
[email protected](
+ "exception_class",
+ [
+ ObjectStoreCreationException,
+ FileFormatRegistrationException,
+ QueryExecutionException,
+ IcebergRegistrationException,
+ ],
+)
+def test_datafusion_exceptions_extend_airflow_exception(exception_class):
+ assert issubclass(exception_class, AirflowException)
