# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
from typing import Any, Iterator, List, MutableMapping, Optional, Dict

import json
import os
import pytest
from unittest.mock import MagicMock, Mock, patch

from azure.ai.agents import AgentsClient
from azure.ai.agents.operations import RunsOperations
from azure.ai.agents.models import (
    CodeInterpreterTool,
    FunctionTool,
    RequiredFunctionToolCall,
    RequiredFunctionToolCallDetails,
    RequiredToolCall,
    RunStatus,
    SubmitToolOutputsAction,
    SubmitToolOutputsDetails,
    ToolSet,
    ToolOutput,
    AgentEventHandler,
    ThreadRun,
)

from user_functions import user_functions


JSON = MutableMapping[str, Any]  # pylint: disable=unsubscriptable-object


def read_file(file_name: str) -> str:
    with open(os.path.join(os.path.dirname(__file__), "assets", f"{file_name}.txt"), "r") as file:
        return file.read()


main_stream_response = read_file("main_stream_response")
fetch_current_datetime_and_weather_stream_response = read_file("fetch_current_datetime_and_weather_stream_response")
send_email_stream_response = read_file("send_email_stream_response")


def convert_to_byte_iterator(main_stream_response: str) -> Iterator[bytes]:
    yield main_stream_response.encode()


def function1():
    return "output from the first agent"


def function2():
    return "output from the second agent"


def function_throw_exception():
    raise ValueError("Just a minute")


def _read_file(file_name: str) -> str:
    with open(os.path.join(os.path.dirname(__file__), "assets", f"{file_name}.txt"), "r") as file:
        return file.read()


def _convert_to_byte_iterator(input: str) -> Iterator[bytes]:
    yield input.encode()


def _get_stream_with_tool_calls() -> Iterator[bytes]:
    fetch_current_datetime_and_weather_stream_response = _read_file(
        "fetch_current_datetime_and_weather_stream_response"
    )
    return _convert_to_byte_iterator(fetch_current_datetime_and_weather_stream_response)


class TestAgentsMock:
    """Tests for agent operations"""

    LOCAL_FN = {function1.__name__: function1, function2.__name__: function2}

    def get_mock_client(self) -> AgentsClient:
        """Return the fake project client"""
        client = AgentsClient(
            endpoint="www.bcac95dd-a1eb-11ef-978f-8c1645fec84b.com",
            credential=MagicMock(),
        )
        client.runs.submit_tool_outputs = MagicMock()
        return client

    def get_toolset(self, file_id: Optional[str], function: Optional[str]) -> Optional[ToolSet]:
        """Get the tool set with given file id and function"""
        if file_id is None or function is None:
            return None
        functions = FunctionTool({function})
        code_interpreter = CodeInterpreterTool(file_ids=[file_id])
        toolset = ToolSet()
        toolset.add(functions)
        toolset.add(code_interpreter)
        return toolset

    def _assert_pipeline_and_reset(self, mock_pipeline_run: MagicMock, tool_set: Optional[ToolSet]) -> None:
        """Check that the pipeline has correct values of tools."""
        mock_pipeline_run.assert_called_once()
        data = json.loads(mock_pipeline_run.call_args_list[0].args[0].body)
        assert isinstance(data, dict), f"Wrong body JSON type expected dict, found {type(data)}"
        if tool_set is not None:
            assert "tool_resources" in data, "tool_resources must be in data"
            assert "tools" in data, "tools must be in data"
            expected_file_id = tool_set.resources.code_interpreter.file_ids[0]
            expected_function_name = tool_set.definitions[0].function.name
            # Check code interpreter file id.
            assert data["tool_resources"], "Tools resources is empty."
            assert "code_interpreter" in data["tool_resources"]
            assert data["tool_resources"]["code_interpreter"], "Code interpreter section is empty."
            assert "file_ids" in data["tool_resources"]["code_interpreter"]
            assert (
                data["tool_resources"]["code_interpreter"]["file_ids"][0] == expected_file_id
            ), f"{expected_file_id[0]=}, but found {data['tool_resources']['code_interpreter']['file_ids']}"
            # Check tools.
            assert data["tools"], "Tools must not be empty"
            assert "function" in data["tools"][0]
            assert "name" in data["tools"][0]["function"]
            assert (
                data["tools"][0]["function"]["name"] == expected_function_name
            ), f"{expected_function_name=}, but encountered {data['tools'][0]['function']['name']}"
        else:
            assert "tool_resources" not in data, "tool_resources must not be in data"
            assert "tools" not in data, "tools must not be in data"
        mock_pipeline_run.reset_mock()

    def _get_agent_json(self, name: str, agent_id: str, tool_set: Optional[ToolSet]) -> Dict[str, Any]:
        """Read in the agent JSON, so that we can assume service returnred it."""
        with open(
            os.path.join(os.path.dirname(__file__), "test_data", "agent.json"),
            "r",
        ) as fp:
            agent_dict: Dict[str, Any] = json.load(fp)
        assert isinstance(agent_dict, dict)
        agent_dict["name"] = name
        agent_dict["id"] = agent_id
        if tool_set is not None:
            agent_dict["tool_resources"] = tool_set.resources.as_dict()
            agent_dict["tools"] = tool_set.definitions
        return agent_dict

    def _get_run(
        self, thread_id: str, tool_set: Optional[ToolSet], add_azure_fn: bool = False, is_complete: bool = False
    ) -> Dict[str, Any]:
        """Return JSON as if we have created the run."""
        with open(
            os.path.join(
                os.path.dirname(__file__),
                "test_data",
                "thread_run.json",
            ),
            "r",
        ) as fp:
            run_dict: Dict[str, Any] = json.load(fp)
        run_dict["id"] = thread_id
        run_dict["assistant_id"] = thread_id[3:]
        assert isinstance(run_dict, dict)
        if is_complete:
            run_dict["status"] = RunStatus.COMPLETED
        tool_calls = []
        definitions = []
        if add_azure_fn:
            tool_calls.append(RequiredToolCall(id="1", type="azure_function"))
            definitions.append(
                {
                    "type": "azure_function",
                    "azure_function": {
                        "function": {
                            "name": "foo",
                            "description": "Get answers from the foo bot.",
                            "parameters": {
                                "type": "object",
                                "properties": {
                                    "query": {"type": "string", "description": "The question to ask."},
                                    "outputqueueuri": {"type": "string", "description": "The full output queue uri."},
                                },
                                "required": ["query"],
                            },
                        },
                        "input_binding": {
                            "type": "storage_queue",
                            "storage_queue": {
                                "queue_service_uri": "https://example.windows.net",
                                "queue_name": "azure-function-foo-input",
                            },
                        },
                        "output_binding": {
                            "type": "storage_queue",
                            "storage_queue": {
                                "queue_service_uri": "https://example.queue.core.windows.net",
                                "queue_name": "azure-function-tool-output",
                            },
                        },
                    },
                }
            )
        if tool_set is not None:
            tool_calls.append(
                RequiredFunctionToolCall(
                    id="0",
                    function=RequiredFunctionToolCallDetails(
                        name=tool_set.definitions[0].function.name,
                        arguments="{}",
                    ),
                )
            )
            definitions.extend(tool_set.definitions)
            run_dict["tool_resources"] = tool_set.resources.as_dict()
        if tool_calls:
            sb = SubmitToolOutputsAction(submit_tool_outputs=SubmitToolOutputsDetails(tool_calls=tool_calls))
            run_dict["required_action"] = sb.as_dict()
            run_dict["tools"] = definitions
        return run_dict

    def _assert_tool_call(self, submit_tool_mock: MagicMock, run_id: str, tool_set: Optional[ToolSet]) -> None:
        """Check that submit_tool_outputs_to_run was called with correct parameters or was not called"""
        if tool_set is not None:
            expected_out = TestAgentsMock.LOCAL_FN[tool_set.definitions[0].function.name]()
            submit_tool_mock.assert_called_once()
            submit_tool_mock.assert_called_with(
                thread_id="some_thread_id",
                run_id=run_id,
                tool_outputs=[{"tool_call_id": "0", "output": expected_out}],
            )
            submit_tool_mock.reset_mock()
        else:
            submit_tool_mock.assert_not_called()

    def _set_toolcalls(
        self, agents_client: AgentsClient, toolset1: Optional[ToolSet], toolset2: Optional[ToolSet]
    ) -> None:
        """Set the tool calls for the agent."""
        max_retry = 3
        if toolset1 and toolset2:
            function_in_toolset1 = set(toolset1.get_tool(tool_type=FunctionTool)._functions.values())
            function_in_toolset2 = set(toolset2.get_tool(tool_type=FunctionTool)._functions.values())
            function_tool = FunctionTool(function_in_toolset1)
            function_tool.add_functions(function_in_toolset2)
            agents_client.enable_auto_function_calls(function_tool, max_retry=max_retry)
        elif toolset1:
            agents_client.enable_auto_function_calls(toolset1, max_retry=max_retry)
        elif toolset2:
            agents_client.enable_auto_function_calls(toolset2, max_retry=max_retry)

    @patch("azure.ai.agents._client.PipelineClient")
    @pytest.mark.parametrize(
        "file_agent_1,file_agent_2",
        [
            ("file_for_agent1", "file_for_agent2"),
            (None, "file_for_agent2"),
            ("file_for_agent1", None),
            (None, None),
        ],
    )
    def test_multiple_agents_create(
        self,
        mock_pipeline_client_gen: MagicMock,
        file_agent_1: Optional[str],
        file_agent_2: Optional[str],
    ) -> None:
        """Test agents can get correct toolset."""
        toolset1 = self.get_toolset(file_agent_1, function1)
        toolset2 = self.get_toolset(file_agent_2, function2)

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.side_effect = [
            self._get_agent_json("first", "123", toolset1),
            self._get_agent_json("second", "456", toolset2),
            self._get_run("run123", toolset1),  # create_run
            self._get_run("run123", toolset1),  # get_run
            self._get_run("run123", toolset1, is_complete=True),  # get_run after resubmitting with tool results
            self._get_run("run456", toolset2),  # create_run
            self._get_run("run456", toolset2),  # get_run
            self._get_run("run456", toolset2, is_complete=True),  # get_run after resubmitting with tool results
            "{}",  # delete agent 1
            "{}",  # delete agent 2
        ]
        mock_pipeline_response = MagicMock()
        mock_pipeline_response.http_response = mock_response
        mock_pipeline = MagicMock()
        mock_pipeline._pipeline.run.return_value = mock_pipeline_response
        mock_pipeline_client_gen.return_value = mock_pipeline
        agents_client = self.get_mock_client()
        with agents_client:
            self._set_toolcalls(agents_client, toolset1, toolset2)
            # Check that pipelines are created as expected.
            agent1 = agents_client.create_agent(
                model="gpt-4-1106-preview",
                name="first",
                instructions="You are a helpful agent",
                toolset=toolset1,
            )
            self._assert_pipeline_and_reset(mock_pipeline._pipeline.run, tool_set=toolset1)

            agent2 = agents_client.create_agent(
                model="gpt-4-1106-preview",
                name="second",
                instructions="You are a helpful agent",
                toolset=toolset2,
            )
            self._assert_pipeline_and_reset(mock_pipeline._pipeline.run, tool_set=toolset2)
            # Check that the new agents are called with correct tool sets.
            agents_client.runs.create_and_process(thread_id="some_thread_id", agent_id=agent1.id, polling_interval=0)
            self._assert_tool_call(agents_client.runs.submit_tool_outputs, "run123", toolset1)

            agents_client.runs.create_and_process(thread_id="some_thread_id", agent_id=agent2.id, polling_interval=0)
            self._assert_tool_call(agents_client.runs.submit_tool_outputs, "run456", toolset2)
            # Check that we cleanup tools after deleting agent.
            agents_client.delete_agent(agent1.id)
            agents_client.delete_agent(agent2.id)

    @patch("azure.ai.agents.operations._operations.RunsOperations.cancel")
    @patch("azure.ai.agents._client.PipelineClient")
    def test_auto_function_calls_retry(
        self,
        mock_pipeline_client_gen: MagicMock,
        mock_cancel_run: MagicMock,
    ) -> None:
        """Test azure function with toolset."""
        toolset = self.get_toolset("file_for_agent1", function_throw_exception)
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.side_effect = [
            self._get_agent_json("first", "123", toolset),
            self._get_run("run2", toolset),  # create_run
            self._get_run("run2", toolset),  # get_run
            self._get_run("run3", toolset),  # get_run
            self._get_run("run4", toolset),  # get_run
            self._get_run("run5", toolset),  # get_run
        ]
        mock_pipeline_response = MagicMock()
        mock_pipeline_response.http_response = mock_response
        mock_pipeline = MagicMock()
        mock_pipeline._pipeline.run.return_value = mock_pipeline_response
        mock_pipeline_client_gen.return_value = mock_pipeline
        agents_client = self.get_mock_client()
        with agents_client:
            # Check that pipelines are created as expected.
            self._set_toolcalls(agents_client, toolset, None)
            agent1 = agents_client.create_agent(
                model="gpt-4-1106-preview",
                name="first",
                instructions="You are a helpful assistant",
                toolset=toolset,
            )
            # Create run with new tool set, which also can be none.
            agents_client.runs.create_and_process(thread_id="some_thread_id", agent_id=agent1.id)
            assert mock_cancel_run.call_count == 1
            assert agents_client.runs.submit_tool_outputs.call_count == 3

    @patch("azure.ai.agents.operations._operations.RunsOperations.cancel")
    @patch("azure.ai.agents._client.PipelineClient")
    def test_auto_function_calls_in_stream(
        self,
        mock_pipeline_client_gen: MagicMock,
        mock_cancel_run: MagicMock,
    ) -> None:
        """Test azure function with toolset."""
        toolset = self.get_toolset("file_for_agent1", function_throw_exception)
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.side_effect = [
            self._get_agent_json("first", "123", toolset),
        ]
        mock_cancel_run.return_value = ThreadRun({"id": "1", "status": RunStatus.CANCELLED})

        mock_pipeline_response = MagicMock()
        mock_pipeline_response.http_response = mock_response
        mock_pipeline = MagicMock()
        mock_pipeline._pipeline.run.return_value = mock_pipeline_response
        mock_pipeline_client_gen.return_value = mock_pipeline
        agents_client = self.get_mock_client()
        with agents_client:
            # Check that pipelines are created as expected.
            self._set_toolcalls(agents_client, toolset, None)
            agent1 = agents_client.create_agent(
                model="gpt-4-1106-preview",
                name="first",
                instructions="You are a helpful assistant",
                toolset=toolset,
            )
            # Create run with new tool set, which also can be none.

            mock_response.iter_bytes.side_effect = [
                _get_stream_with_tool_calls(),  # create_run
                _get_stream_with_tool_calls(),  # submit_tool_outputs_to_run
                _get_stream_with_tool_calls(),  # submit_tool_outputs_to_run
                _get_stream_with_tool_calls(),  # submit_tool_outputs_to_run
            ]

            event_handler = AgentEventHandler()
            with agents_client.runs.stream(
                thread_id="some_thread_id", agent_id=agent1.id, event_handler=event_handler
            ) as stream:
                stream.until_done()
            assert mock_cancel_run.call_count == 1
            assert event_handler.current_retry == 4

    @patch("azure.ai.agents._client.PipelineClient")
    @pytest.mark.parametrize(
        "file_agent_1,file_agent_2",
        [
            ("file_for_agent1", "file_for_agent2"),
            (None, "file_for_agent2"),
            ("file_for_agent1", None),
            (None, None),
        ],
    )
    def test_update_agent_tools(
        self,
        mock_pipeline_client_gen: MagicMock,
        file_agent_1: Optional[str],
        file_agent_2: Optional[str],
    ) -> None:
        """Test that tools are properly updated."""
        toolset1 = self.get_toolset(file_agent_1, function1)
        toolset2 = self.get_toolset(file_agent_2, function2)
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.side_effect = [
            self._get_agent_json("first", "123", toolset1),
            self._get_agent_json("first", "123", toolset2),
        ]
        mock_pipeline_response = MagicMock()
        mock_pipeline_response.http_response = mock_response
        mock_pipeline = MagicMock()
        mock_pipeline._pipeline.run.return_value = mock_pipeline_response
        mock_pipeline_client_gen.return_value = mock_pipeline
        agents_client = self.get_mock_client()
        with agents_client:
            # Check that pipelines are created as expected.
            agent1 = agents_client.create_agent(
                model="gpt-4-1106-preview",
                name="first",
                instructions="You are a helpful agent",
                toolset=toolset1,
            )
            agent1 = agents_client.update_agent(agent1.id, toolset=toolset2)
            if toolset2 is None:
                assert agent1.tools == None
            else:
                assert agent1.tools[0].function.name == function2.__name__

    @patch("azure.ai.agents._client.PipelineClient")
    @pytest.mark.parametrize(
        "file_agent_1,file_agent_2",
        [
            ("file_for_agent1", "file_for_agent2"),
            (None, "file_for_agent2"),
            ("file_for_agent1", None),
            (None, None),
        ],
    )
    def test_create_run_tools_override(
        self,
        mock_pipeline_client_gen: MagicMock,
        file_agent_1: Optional[str],
        file_agent_2: Optional[str],
    ) -> None:
        """Test that if user have set tool set in create create_and_process_run method, that tools are used."""
        toolset1 = self.get_toolset(file_agent_1, function1)
        toolset2 = self.get_toolset(file_agent_2, function2)
        mock_response = MagicMock()
        mock_response.status_code = 200
        side_effect = [self._get_agent_json("first", "123", toolset1)]
        if toolset1 is not None or toolset2 is not None:
            toolset = toolset2 if toolset2 is not None else toolset1
            side_effect.append(self._get_run("run123", toolset))  # create_run
            side_effect.append(self._get_run("run123", toolset))  # get_run
            side_effect.append(
                self._get_run("run123", toolset, is_complete=True)
            )  # get_run after resubmitting with tool results
        else:
            side_effect.append(
                self._get_run("run123", None, is_complete=True)
            )  # Run must be marked as completed in this case.
        mock_response.json.side_effect = side_effect
        mock_pipeline_response = MagicMock()
        mock_pipeline_response.http_response = mock_response
        mock_pipeline = MagicMock()
        mock_pipeline._pipeline.run.return_value = mock_pipeline_response
        mock_pipeline_client_gen.return_value = mock_pipeline
        agents_client = self.get_mock_client()
        with agents_client:
            # Check that pipelines are created as expected.
            self._set_toolcalls(agents_client, toolset1, toolset2)
            agent1 = agents_client.create_agent(
                model="gpt-4-1106-preview",
                name="first",
                instructions="You are a helpful agent",
                toolset=toolset1,
            )
            self._assert_pipeline_and_reset(mock_pipeline._pipeline.run, tool_set=toolset1)

            # Create run with new tool set, which also can be none.
            agents_client.runs.create_and_process(
                thread_id="some_thread_id", agent_id=agent1.id, toolset=toolset2, polling_interval=0
            )
            if toolset2 is not None:
                self._assert_tool_call(agents_client.runs.submit_tool_outputs, "run123", toolset2)
            else:
                self._assert_tool_call(agents_client.runs.submit_tool_outputs, "run123", toolset1)

    @patch("azure.ai.agents._client.PipelineClient")
    @pytest.mark.parametrize(
        "file_agent_1,add_azure_fn",
        [
            ("file_for_agent1", True),
            (None, True),
            ("file_for_agent1", False),
            (None, False),
        ],
    )
    def test_with_azure_function(
        self,
        mock_pipeline_client_gen: MagicMock,
        file_agent_1: Optional[str],
        add_azure_fn: bool,
    ) -> None:
        """Test azure function with toolset."""
        toolset = self.get_toolset(file_agent_1, function1)
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.side_effect = [
            self._get_agent_json("first", "123", toolset),
            self._get_run("run123", toolset, add_azure_fn=add_azure_fn),  # create_run
            self._get_run("run123", toolset, add_azure_fn=add_azure_fn),  # get_run
            self._get_run(
                "run123", toolset, add_azure_fn=add_azure_fn, is_complete=True
            ),  # get_run after resubmitting with tool results
        ]
        mock_pipeline_response = MagicMock()
        mock_pipeline_response.http_response = mock_response
        mock_pipeline = MagicMock()
        mock_pipeline._pipeline.run.return_value = mock_pipeline_response
        mock_pipeline_client_gen.return_value = mock_pipeline
        agents_client = self.get_mock_client()
        with agents_client:
            # Check that pipelines are created as expected.
            self._set_toolcalls(agents_client, toolset, None)
            agent1 = agents_client.create_agent(
                model="gpt-4-1106-preview",
                name="first",
                instructions="You are a helpful agent",
                toolset=toolset,
            )
            # Create run with new tool set, which also can be none.
            agents_client.runs.create_and_process(thread_id="some_thread_id", agent_id=agent1.id, polling_interval=0)
            self._assert_tool_call(agents_client.runs.submit_tool_outputs, "run123", toolset)

    def _assert_stream_call(self, submit_tool_mock: MagicMock, run_id: str, tool_set: Optional[ToolSet]) -> None:
        """Assert that stream has received the correct values."""
        if tool_set is not None:
            expected_out = TestAgentsMock.LOCAL_FN[tool_set.definitions[0].function.name]()
            submit_tool_mock.assert_called_once()
            submit_tool_mock.assert_called_with(
                thread_id="some_thread_id",
                run_id=run_id,
                tool_outputs=[{"tool_call_id": "0", "output": expected_out}],
                event_handler=None,
            )
            submit_tool_mock.reset_mock()
        else:
            submit_tool_mock.assert_not_called()

    @patch("azure.ai.agents._client.PipelineClient")
    @pytest.mark.skip("Recordings not yet available")
    @pytest.mark.parametrize(
        "file_agent_1,add_azure_fn",
        [
            ("file_for_agent1", True),
            (None, True),
            ("file_for_agent1", False),
            (None, False),
        ],
    )
    def test_handle_submit_tool_outputs(
        self,
        mock_pipeline_client_gen: MagicMock,
        file_agent_1: Optional[str],
        add_azure_fn: bool,
    ) -> None:
        """Test handling of stream tools response."""
        toolset = self.get_toolset(file_agent_1, function1)
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.side_effect = [
            self._get_agent_json("first", "123", toolset),
            self._get_run("run123", toolset, add_azure_fn=add_azure_fn),  # create_run
            self._get_run("run123", toolset, add_azure_fn=add_azure_fn),  # get_run
            self._get_run(
                "run123", toolset, add_azure_fn=add_azure_fn, is_complete=True
            ),  # get_run after resubmitting with tool results
        ]
        mock_pipeline_response = MagicMock()
        mock_pipeline_response.http_response = mock_response
        mock_pipeline = MagicMock()
        mock_pipeline._pipeline.run.return_value = mock_pipeline_response
        mock_pipeline_client_gen.return_value = mock_pipeline
        agents_client = self.get_mock_client()
        with agents_client:
            # Check that pipelines are created as expected.
            self._set_toolcalls(agents_client, toolset, None)
            agent1 = agents_client.create_agent(
                model="gpt-4-1106-preview",
                name="first",
                instructions="You are a helpful agent",
                toolset=toolset,
            )
            # Create run with new tool set, which also can be none.
            run = agents_client.runs.create_and_process(
                thread_id="some_thread_id", agent_id=agent1.id, polling_interval=0
            )
            self._assert_tool_call(agents_client.submit_tool_outputs_to_run, "run123", toolset)
            agents_client._handle_submit_tool_outputs(run)
            self._assert_stream_call(agents_client.submit_tool_outputs_to_stream, "run123", toolset)


class TestIntegrationAgentsMock:

    def submit_tool_outputs(
        self, thread_id: str, run_id: str, *, tool_outputs: List[ToolOutput], stream_parameter: bool, stream: bool
    ) -> Iterator[bytes]:
        assert thread_id == "thread_01"
        assert run_id == "run_01"
        assert stream_parameter == True
        assert stream == True
        if (
            len(tool_outputs) == 2
            and tool_outputs[0]["tool_call_id"] == "call_01"
            and tool_outputs[1]["tool_call_id"] == "call_02"
        ):
            return convert_to_byte_iterator(fetch_current_datetime_and_weather_stream_response)
        elif len(tool_outputs) == 1 and tool_outputs[0]["tool_call_id"] == "call_03":
            return convert_to_byte_iterator(send_email_stream_response)
        raise ValueError("Unexpected tool outputs")

    @patch(
        "azure.ai.agents.operations._operations.RunsOperations.create",
        return_value=convert_to_byte_iterator(main_stream_response),
    )
    @patch("azure.ai.agents.AgentsClient.__init__", return_value=None)
    @patch(
        "azure.ai.agents.operations._operations.RunsOperations.submit_tool_outputs",
    )
    def test_create_stream_with_tool_calls(self, mock_submit_tool_outputs_to_run: Mock, *args):
        mock_submit_tool_outputs_to_run.side_effect = self.submit_tool_outputs
        functions = FunctionTool(user_functions)
        toolset = ToolSet()
        toolset.add(functions)

        operation = AgentsClient(
            endpoint="www.bcac95dd-a1eb-11ef-978f-8c1645fec84b.com",
            credential=MagicMock(),
        )
        operation.runs = RunsOperations(MagicMock(), MagicMock(), MagicMock(), MagicMock())
        operation.enable_auto_function_calls(toolset)
        count = 0

        with operation.runs.stream(thread_id="thread_id", agent_id="asst_01") as stream:
            for _ in stream:
                count += 1
        assert count == (
            main_stream_response.count("event:")
            + fetch_current_datetime_and_weather_stream_response.count("event:")
            + send_email_stream_response.count("event:")
        )