Lee-W commented on code in PR #74626:
URL: https://github.com/apache/airflow/pull/74626#discussion_r4262234633
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -351,6 +379,68 @@ def test_builds_the_connection_hook_with_hook_params(self,
registry):
assert hook.model_id == "openai:gpt-5"
assert hook.fallback_conn_ids == []
+ def test_rejects_a_connection_of_another_type_before_dispatching(self,
registry):
+ registry.add("llm", conn_type="openai")
+ registry.register_hook("openai", qualname(_OtherProviderHook))
+ expected_types = ["pydanticai", "pydanticai_azure",
"pydanticai_bedrock", "pydanticai_vertex"]
+
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ f"Connection 'llm' has type 'openai', but a Pydantic AI hook
needs one of {expected_types}"
+ ),
+ ):
+ PydanticAIHook.get_hook("llm", hook_params={"model_id":
"openai:gpt-5"})
+
+ Connection.get_hook.assert_not_called()
+
+ def test_accepts_a_subclass_from_a_module_not_imported_yet(self, registry,
tmp_path, monkeypatch):
+ """Another provider's subclass (e.g. Snowflake's) is only imported
when its type is looked up."""
+ module_name = "pydanticai_subclass_not_imported_yet"
+ (tmp_path / f"{module_name}.py").write_text(
+ "from airflow.providers.common.ai.hooks.pydantic_ai import
PydanticAIHook\n\n\n"
+ "class PydanticAIElsewhereHook(PydanticAIHook):\n"
+ " conn_type = 'pydanticai_elsewhere'\n"
+ )
+ monkeypatch.syspath_prepend(tmp_path)
+ registry.add("llm", conn_type="pydanticai_elsewhere")
Review Comment:
The subclass in the lazy-import test is now registered under
`snowflake_llm`, which has no `pydanticai` prefix, and the test asserts
`module_name in sys.modules` after the call. A `startswith("pydanticai")` check
would now fail it.
--
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]
