@@ -743,15 +743,31 @@ def test_get_provider_kwargs_api_key_gla_mode(self):
743743 )
744744 assert result ["api_key" ] == "gla-key"
745745
746- def test_get_provider_kwargs_vertexai_flag (self ):
747- """vertexai bool is forwarded and coerced to bool."""
746+ @pytest .mark .parametrize ("vertexai_value" , [True , False ])
747+ def test_get_provider_kwargs_vertexai_flag_is_not_forwarded (self , vertexai_value ):
748+ """The ``vertexai`` extra field must never reach the provider constructor.
749+
750+ Neither ``GoogleProvider`` nor ``GoogleCloudProvider`` in current pydantic-ai
751+ accept a ``vertexai`` kwarg (pydantic/pydantic-ai#5336 hardcoded it inside
752+ ``GoogleCloudProvider`` instead). Forwarding it raises ``TypeError``, which the
753+ base hook's fallback then swallows by dropping every other kwarg -- silently
754+ re-resolving credentials from the environment. Regression test for that bug.
755+ """
748756 hook = PydanticAIVertexHook .__new__ (PydanticAIVertexHook )
749757 result = hook ._get_provider_kwargs (
750758 None ,
751759 None ,
752- {"model" : "google-vertex:gemini-2.0-flash" , "api_key" : "key" , "vertexai" : True },
760+ {
761+ "model" : "google-vertex:gemini-2.0-flash" ,
762+ "project" : "my-project" ,
763+ "location" : "us-central1" ,
764+ "vertexai" : vertexai_value ,
765+ },
753766 )
754- assert result ["vertexai" ] is True
767+ assert "vertexai" not in result
768+ # The other credential kwargs must still go through untouched.
769+ assert result ["project" ] == "my-project"
770+ assert result ["location" ] == "us-central1"
755771
756772 def test_get_provider_kwargs_service_account_info_loads_credentials (self ):
757773 """service_account_info dict is loaded into a Credentials object."""
@@ -833,3 +849,67 @@ def test_get_conn_uses_explicit_project(self, mock_infer_provider_class, mock_in
833849 factory = mock_infer_model .call_args [1 ]["provider_factory" ]
834850 factory ("google-vertex" )
835851 mock_provider_cls .assert_called_with (project = "my-project" , location = "europe-west4" )
852+
853+ @patch ("airflow.providers.common.ai.hooks.pydantic_ai.infer_model" , autospec = True )
854+ @patch ("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider" , autospec = True )
855+ @patch ("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class" , autospec = True )
856+ def test_get_conn_vertexai_flag_does_not_trigger_typeerror_fallback (
857+ self , mock_infer_provider_class , mock_infer_provider , mock_infer_model
858+ ):
859+ """Setting ``vertexai`` must not push ``get_conn`` onto the ``except TypeError``
860+ fallback path, which would silently discard project/location/credentials.
861+
862+ The stand-in below has the exact keyword-only signature of the real
863+ ``GoogleCloudProvider.__init__`` (verified against the installed pydantic-ai) so
864+ it raises ``TypeError`` on an unexpected ``vertexai`` kwarg exactly like the real
865+ class would -- the real class itself needs the optional ``google-genai``
866+ dependency, which isn't part of this provider's test environment.
867+ """
868+
869+ class FakeGoogleCloudProvider :
870+ def __init__ (
871+ self ,
872+ * ,
873+ api_key = None ,
874+ credentials = None ,
875+ project = None ,
876+ location = None ,
877+ client = None ,
878+ http_client = None ,
879+ base_url = None ,
880+ retry_options = None ,
881+ ):
882+ self .kwargs = {
883+ "api_key" : api_key ,
884+ "credentials" : credentials ,
885+ "project" : project ,
886+ "location" : location ,
887+ }
888+
889+ mock_infer_model .return_value = MagicMock (spec = Model )
890+ mock_infer_provider_class .return_value = FakeGoogleCloudProvider
891+
892+ hook = PydanticAIVertexHook (llm_conn_id = "vertex_test" )
893+ conn = Connection (
894+ conn_id = "vertex_test" ,
895+ conn_type = "pydanticai-vertex" ,
896+ extra = json .dumps (
897+ {
898+ "model" : "google-cloud:gemini-2.0-flash" ,
899+ "project" : "my-project" ,
900+ "location" : "us-central1" ,
901+ "vertexai" : True ,
902+ }
903+ ),
904+ )
905+ with patch .object (hook , "get_connection" , return_value = conn ):
906+ hook .get_conn ()
907+
908+ factory = mock_infer_model .call_args [1 ]["provider_factory" ]
909+ provider = factory ("google-cloud" )
910+
911+ assert isinstance (provider , FakeGoogleCloudProvider )
912+ assert provider .kwargs ["project" ] == "my-project"
913+ assert provider .kwargs ["location" ] == "us-central1"
914+ # The TypeError fallback must never have been reached.
915+ mock_infer_provider .assert_not_called ()
0 commit comments