Eng 1095 ai xplain sdk create functional tests updating tools and agents

tests/functional/agent/agent_functional_test.py

@@ -163,3 +163,46 @@ def test_delete_agent_in_use(delete_agents_and_team_agents):
     with pytest.raises(Exception) as exc_info:
         agent.delete()
     assert str(exc_info.value) == "Agent Deletion Error (HTTP 403): err.agent_is_in_use."
+
+
+def test_update_tools_of_agent(run_input_map, delete_agents_and_team_agents):
+    assert delete_agents_and_team_agents
+
+    agent = AgentFactory.create(
+        name=run_input_map["agent_name"], description=run_input_map["agent_name"], llm_id=run_input_map["llm_id"]
+    )
+    assert agent is not None
+    assert agent.status == AssetStatus.DRAFT
+    assert len(agent.tools) == 0 
+
+    tools = []
+    if "model_tools" in run_input_map:
+        for tool in run_input_map["model_tools"]:
+            tool_ = copy.copy(tool)
+            for supplier in Supplier:
+                if tool["supplier"] is not None and tool["supplier"].lower() in [
+                    supplier.value["code"].lower(),
+                    supplier.value["name"].lower(),
+                ]:
+                    tool_["supplier"] = supplier
+                    break
+            tools.append(AgentFactory.create_model_tool(**tool_))
+
+    if "pipeline_tools" in run_input_map:
+        for tool in run_input_map["pipeline_tools"]:
+            tools.append(AgentFactory.create_pipeline_tool(pipeline=tool["pipeline_id"], description=tool["description"]))
+
+    agent.tools = tools
+    agent.update()
+
+    agent = AgentFactory.get(agent.id)
+    assert len(agent.tools) == len(tools)  
+
+    removed_tool = agent.tools.pop()
+    agent.update()
+
+    agent = AgentFactory.get(agent.id)
+    assert len(agent.tools) == len(tools) - 1 
+    assert removed_tool not in agent.tools 
+
+    agent.delete()

tests/functional/team_agent/team_agent_functional_test.py

@@ -160,3 +160,62 @@ def test_fail_non_existent_llm():
             tools=[AgentFactory.create_model_tool(function=Function.TRANSLATION)],
         )
     assert str(exc_info.value) == "Large Language Model with ID 'non_existent_llm' not found."
+
+def test_add_remove_agents_from_team_agent(run_input_map, delete_agents_and_team_agents):
+    assert delete_agents_and_team_agents
+
+    agents = []
+    for agent in run_input_map["agents"]:
+        tools = []
+        if "model_tools" in agent:
+            for tool in agent["model_tools"]:
+                tool_ = copy(tool)
+                for supplier in Supplier:
+                    if tool["supplier"] is not None and tool["supplier"].lower() in [
+                        supplier.value["code"].lower(),
+                        supplier.value["name"].lower(),
+                    ]:
+                        tool_["supplier"] = supplier
+                        break
+                tools.append(AgentFactory.create_model_tool(**tool_))
+        if "pipeline_tools" in agent:
+            for tool in agent["pipeline_tools"]:
+                tools.append(AgentFactory.create_pipeline_tool(pipeline=tool["pipeline_id"], description=tool["description"]))
+
+        agent = AgentFactory.create(
+            name=agent["agent_name"], description=agent["agent_name"], llm_id=agent["llm_id"], tools=tools
+        )
+        agents.append(agent)
+
+    team_agent = TeamAgentFactory.create(
+        name=run_input_map["team_agent_name"],
+        agents=agents,
+        description=run_input_map["team_agent_name"],
+        llm_id=run_input_map["llm_id"],
+        use_mentalist_and_inspector=True,
+    )
+
+    assert team_agent is not None
+    assert team_agent.status == AssetStatus.DRAFT
+
+    new_agent = AgentFactory.create(
+        name="New Agent",
+        description="Agent added to team",
+        llm_id=run_input_map["llm_id"],
+    )
+    team_agent.agents.append(new_agent)
+    team_agent.update()
+
+    team_agent = TeamAgentFactory.get(team_agent.id)
+    assert new_agent.id in [agent.id for agent in team_agent.agents]
+    assert len(team_agent.agents) == len(agents) + 1
+
+    removed_agent = team_agent.agents.pop(0)  
+    team_agent.update()
+
+    team_agent = TeamAgentFactory.get(team_agent.id)
+    assert removed_agent.id not in [agent.id for agent in team_agent.agents]
+    assert len(team_agent.agents) == len(agents) 
+
+    team_agent.delete()
+    new_agent.delete()