Fixup tests

Author: gopidesupavan

providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py

@@ -753,31 +753,30 @@ def test_instrument_set_when_settings_returned(self, mock_settings):
         sentinel = MagicMock(name="InstrumentationSettings")
         mock_settings.return_value = sentinel
         hook = self._hook()
-        with patch.object(hook, "get_conn", return_value=TestModel()):
-            agent = hook.create_agent(instructions="hi")
+        with patch.object(hook, "get_model", return_value=TestModel()):
+            agent = hook.create_agent(AgentRunRequest(prompt="test", instructions="hi"))
 
         assert agent.instrument is sentinel
 
     @patch("airflow.providers.common.ai.hooks.pydantic_ai.genai_instrumentation_settings")
     def test_no_instrument_when_settings_none(self, mock_settings):
         mock_settings.return_value = None
         hook = self._hook()
-        with patch.object(hook, "get_conn", return_value=TestModel()):
-            agent = hook.create_agent(instructions="hi")
+        with patch.object(hook, "get_model", return_value=TestModel()):
+            agent = hook.create_agent(AgentRunRequest(prompt="test", instructions="hi"))
 
         mock_settings.assert_called_once()
         assert agent.instrument is None
 
     @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True)
     @patch("airflow.providers.common.ai.hooks.pydantic_ai.genai_instrumentation_settings")
-    @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True)
-    def test_caller_instrument_short_circuits(self, mock_infer_model, mock_settings, mock_agent_cls):
-        """A caller that passes its own ``instrument`` wins; we don't override it."""
-        mock_infer_model.return_value = MagicMock(spec=Model)
+    def test_caller_instrument_short_circuits(self, mock_settings, mock_agent_cls):
+        """A caller that passes its own ``instrument`` via agent_params wins; we don't override it."""
         hook = self._hook()
-        conn = Connection(conn_id="test_conn", conn_type="pydanticai")
-        with patch.object(hook, "get_connection", return_value=conn):
-            hook.create_agent(instructions="hi", instrument=False)
+        with patch.object(hook, "get_model", return_value=TestModel()):
+            hook.create_agent(
+                AgentRunRequest(prompt="test", instructions="hi", agent_params={"instrument": False})
+            )
 
         mock_settings.assert_not_called()