diff --git a/tests/utilities/test_planning_auth/test_planning_auth_inheritance.py b/tests/utilities/test_planning_auth/test_planning_auth_inheritance.py index e264c74d3..b38cd8489 100644 --- a/tests/utilities/test_planning_auth/test_planning_auth_inheritance.py +++ b/tests/utilities/test_planning_auth/test_planning_auth_inheritance.py @@ -1,39 +1,25 @@ -from crewai import Agent, Task -from crewai.llm import LLM +from unittest.mock import MagicMock +from crewai import Task from crewai.utilities.planning_handler import CrewPlanner def test_planning_llm_inherits_auth_params(): """Test that planning LLM inherits authentication parameters from agent LLM.""" - custom_llm = LLM( - model="custom-model", - base_url="https://api.custom-provider.com/v1", - api_key="fake-api-key", - api_version="2023-05-15", - organization="custom-org" - ) - - agent = Agent( - role="Test Agent", - goal="Test Goal", - backstory="Test Backstory", - llm=custom_llm - ) + mock_llm = MagicMock() + mock_llm.base_url = "https://api.custom-provider.com/v1" + mock_llm.api_version = "2023-05-15" task = Task( description="Test Task", - expected_output="Test Output", - agent=agent + expected_output="Test Output" ) planner = CrewPlanner( tasks=[task], planning_agent_llm=None, # This should trigger the inheritance logic - agent_llm=custom_llm + agent_llm=mock_llm ) assert hasattr(planner, 'planning_agent_llm') assert hasattr(planner.planning_agent_llm, 'base_url') assert planner.planning_agent_llm.base_url == "https://api.custom-provider.com/v1" - assert planner.planning_agent_llm.api_key == "fake-api-key" assert planner.planning_agent_llm.api_version == "2023-05-15" - # organization is not directly accessible as an attribute