Lee-W commented on code in PR #74626:
URL: https://github.com/apache/airflow/pull/74626#discussion_r4262237788
##########
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")
+ registry.register_hook("pydanticai_elsewhere",
f"{module_name}.PydanticAIElsewhereHook")
+ assert module_name not in sys.modules
+
+ try:
+ PydanticAIHook.get_hook("llm", hook_params={"model_id":
"openai:gpt-5"})
+ finally:
+ sys.modules.pop(module_name, None)
+
+ Connection.get_hook.assert_called_once_with(
+ registry.conns["llm"], hook_params={"model_id": "openai:gpt-5"}
+ )
+
+ def
test_leaves_an_unregistered_connection_type_to_connection_get_hook(self,
registry):
+ """``Connection.get_hook`` raises for it with a hint about the '-'/'_'
spelling."""
+ registry.add("llm", conn_type="pydanticai-azure")
+
+ PydanticAIHook.get_hook("llm")
+
+ Connection.get_hook.assert_called_once_with(registry.conns["llm"],
hook_params=None)
+
+ @pytest.mark.parametrize(
+ ("conn_type", "hook_class"),
+ [
+ ("pydanticai", PydanticAIHook),
+ ("pydanticai_azure", PydanticAIAzureHook),
+ ("pydanticai_bedrock", PydanticAIBedrockHook),
+ ("pydanticai_vertex", PydanticAIVertexHook),
+ ],
+ )
+ def test_accepts_every_pydanticai_connection_type(self, registry,
conn_type, hook_class):
+ registry.add("llm", conn_type=conn_type, hook_class=hook_class)
+
+ hook = PydanticAIHook.get_hook("llm")
+
+ assert type(hook) is hook_class
Review Comment:
`test_accepts_every_pydanticai_connection_type` now checks
`Connection.get_hook.assert_called_once_with(registry.conns["llm"],
hook_params=None)`.
It is parametrized over the `pydanticai*` entries in
`get_provider_info()["connection-types"]`, so it follows provider.yaml.
--
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]
