ParagEkbote commented on code in PR #76317: URL: https://github.com/apache/airflow/pull/76317#discussion_r4496181261
########## providers/huggingface/tests/unit/huggingface/hooks/test_huggingface_hook.py: ########## @@ -0,0 +1,163 @@ +# 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. +"""Connection-backed authentication tests for HuggingFaceHook.""" +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from airflow.exceptions import AirflowNotFoundException + +from airflow.providers.huggingface.hooks.huggingface import HuggingFaceHook +from huggingface_hub import GitCommitInfo + + +def git_commit_info(commit_id, title, message=""): + return GitCommitInfo( + commit_id=commit_id, + authors=["unit-test"], + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + title=title, + message=message, + formatted_title=None, + formatted_message=None, + ) + + [email protected] +def connection_env(monkeypatch): + # Airflow resolves this as a real Connection, without a metadata DB fixture. + monkeypatch.setenv( + "AIRFLOW_CONN_HUGGINGFACE_TEST", + '{"conn_type":"huggingface","password":"test-secret"}', + ) + + +def test_connection_resolution_uses_password_only(connection_env): + hook = HuggingFaceHook("huggingface_test") + assert hook.get_token() == "test-secret" + + +def test_get_conn_caches_autospecced_hf_api(connection_env): + with patch("airflow.providers.huggingface.hooks.huggingface.HfApi", autospec=True) as api: + hook = HuggingFaceHook("huggingface_test") + first = hook.get_conn() + assert hook.get_conn() is first + api.assert_called_once_with(token="test-secret") + + +def test_connection_test_success_and_failure(connection_env): + with patch("airflow.providers.huggingface.hooks.huggingface.HfApi", autospec=True) as api: + hook = HuggingFaceHook("huggingface_test") + assert hook.test_connection() == (True, "Connection successfully tested") + api.return_value.whoami.assert_called_once_with() + api.return_value.whoami.side_effect = ValueError("invalid token") + assert hook.test_connection() == (False, "invalid token") + + +def test_no_connection_defers_to_huggingface_hub_fallback(monkeypatch): + monkeypatch.setenv("HF_TOKEN", "ambient-token") + with patch("airflow.providers.huggingface.hooks.huggingface.HfApi", autospec=True) as api: + hook = HuggingFaceHook(None) + hook.get_conn() + api.assert_called_once_with(token=None) + + +def test_default_missing_connection_allows_sdk_fallback(monkeypatch): + monkeypatch.delenv("AIRFLOW_CONN_HUGGINGFACE_DEFAULT", raising=False) + with patch.object(HuggingFaceHook, "get_connection", side_effect=AirflowNotFoundException("missing")): + assert HuggingFaceHook().get_token() is None + + +def test_explicit_invalid_connection_fails(monkeypatch): + with patch.object(HuggingFaceHook, "get_connection", side_effect=AirflowNotFoundException("missing")): + with pytest.raises(AirflowNotFoundException): + HuggingFaceHook("explicit_bad_id").get_token() + + [email protected]("error_kind", ["value", "gated"]) +def test_load_dataset_passes_token_and_propagates_library_exception(connection_env, error_kind): + from airflow.providers.huggingface.hooks import huggingface + import httpx + + if error_kind == "gated": + from huggingface_hub.utils import GatedRepoError + + error = GatedRepoError( + "gated", response=httpx.Response(403, request=httpx.Request("GET", "https://huggingface.co";)) Review Comment: Modified the test as per the comments. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
