Repository: deepset-ai/haystack

Issue:
Deserialize - AzureOpenAIGenerator
**Describe the bug**
I am using `haystack-ai==2.1.1`

I create prettry simple pipeline which only use prompt builder than it pass it to azure open ai and serialize it, I am trying to load it from another module but I am failling due to the attached trace back.

you can find a notebook which I am using for creating the pipeline also attached.

**Error message**
```python
Traceback (most recent call last):
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/haystack/core/pipeline/pipeline.py", line 172, in from_dict
    raise PipelineError(
haystack.core.errors.PipelineError: Successfully imported module haystack.components.generators but can't find it in the component registry.This is unexpected and most likely a bug.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/uvicorn/protocols/http/httptools_impl.py", line 411, in run_asgi
    result = await app(  # type: ignore[func-returns-value]
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/uvicorn/middleware/proxy_headers.py", line 69, in __call__
    return await self.app(scope, receive, send)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/fastapi/applications.py", line 1054, in __call__
    await super().__call__(scope, receive, send)
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/applications.py", line 123, in __call__
    await self.middleware_stack(scope, receive, send)
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/middleware/errors.py", line 186, in __call__
    raise exc
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/middleware/errors.py", line 164, in __call__
    await self.app(scope, receive, _send)
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/middleware/exceptions.py", line 65, in __call__
    await wrap_app_handling_exceptions(self.app, conn)(scope, receive, send)
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/_exception_handler.py", line 64, in wrapped_app
    raise exc
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/_exception_handler.py", line 53, in wrapped_app
    await app(scope, receive, sender)
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/routing.py", line 756, in __call__
    await self.middleware_stack(scope, receive, send)
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/routing.py", line 776, in app
    await route.handle(scope, receive, send)
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/routing.py", line 297, in handle
    await self.app(scope, receive, send)
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/routing.py", line 77, in app
    await wrap_app_handling_exceptions(app, request)(scope, receive, send)
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/_exception_handler.py", line 64, in wrapped_app
    raise exc
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/_exception_handler.py", line 53, in wrapped_app
    await app(scope, receive, sender)
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/routing.py", line 72, in app
    response = await func(request)
               ^^^^^^^^^^^^^^^^^^^
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/fastapi/routing.py", line 278, in app
    raw_response = await run_endpoint_function(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/fastapi/routing.py", line 193, in run_endpoint_function
    return await run_in_threadpool(dependant.call, **values)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/starlette/concurrency.py", line 42, in run_in_threadpool
    return await anyio.to_thread.run_sync(func, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/anyio/to_thread.py", line 56, in run_sync
    return await get_async_backend().run_sync_in_worker_thread(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/anyio/_backends/_asyncio.py", line 2144, in run_sync_in_worker_thread
    return await future
           ^^^^^^^^^^^^
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/anyio/_backends/_asyncio.py", line 851, in run
    result = context.run(func, *args)
             ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/grabin/projects/netnemo/netguru-service/netguru_service/api/v1/chat.py", line 26, in chat
    pipeline.loads(yaml_pipeline)
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/haystack/core/pipeline/pipeline.py", line 234, in loads
    return cls.from_dict(marshaller.unmarshal(data), callbacks)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/grabin/projects/netnemo/netguru-service/.venv/lib/python3.12/site-packages/haystack/core/pipeline/pipeline.py", line 177, in from_dict
    raise PipelineError(f"Component '{component_data['type']}' not imported.") from e
haystack.core.errors.PipelineError: Component 'haystack.components.generators.AzureOpenAIGenerator' not imported.
```

**Expected behavior**
I expected the 

**Additional context**
Add any other context about the problem here, like document types / preprocessing steps / settings of reader etc.

**To Reproduce**
Steps to reproduce the behavior `AzureOpenAIGenerator` to be loaded correctly.

**FAQ Check**
- [x] Have you had a look at [our new FAQ page](https://docs.haystack.deepset.ai/docs/faq)? yes

**System:**
 - OS: Mac
 - GPU/CPU: M3 pro
 - Haystack version (commit or version number): 2.1.1
 - DocumentStore:
 - Reader:
 - Retriever:


Discussion:
coveralls: ## Pull Request Test Coverage Report for [Build 9083924090](https://coveralls.io/builds/67486271)


### Details

* **0** of **0**   changed or added relevant lines in **0** files are covered.
* **4** unchanged lines in **2** files lost coverage.
* Overall coverage increased (+**0.1%**) to **90.575%**

---


|  Files with Coverage Reduction | New Missed Lines | % |
| :-----|--------------|--: |
| [components/generators/azure.py](https://coveralls.io/builds/67486271/source?filename=components%2Fgenerators%2Fazure.py#L109) | 2 | 92.11% |
| [components/generators/chat/azure.py](https://coveralls.io/builds/67486271/source?filename=components%2Fgenerators%2Fchat%2Fazure.py#L126) | 2 | 91.89% |
<!-- | **Total:** | **4** |  | -->

|  Totals | [![Coverage Status](https://coveralls.io/builds/67486271/badge)](https://coveralls.io/builds/67486271) |
| :-- | --: |
| Change from base [Build 9081766626](https://coveralls.io/builds/67482371): |  0.1% |
| Covered Lines: | 6554 |
| Relevant Lines: | 7236 |

---
##### 💛  - [Coveralls](https://coveralls.io)


Patch:
diff --git a/haystack/components/generators/azure.py b/haystack/components/generators/azure.py
index 388b9eb9..2c432d82 100644
--- a/haystack/components/generators/azure.py
+++ b/haystack/components/generators/azure.py
@@ -8,7 +8,7 @@ from typing import Any, Callable, Dict, Optional
 # pylint: disable=import-error
 from openai.lib.azure import AzureOpenAI
 
-from haystack import default_from_dict, default_to_dict, logging
+from haystack import component, default_from_dict, default_to_dict, logging
 from haystack.components.generators import OpenAIGenerator
 from haystack.dataclasses import StreamingChunk
 from haystack.utils import Secret, deserialize_callable, deserialize_secrets_inplace, serialize_callable
@@ -16,6 +16,7 @@ from haystack.utils import Secret, deserialize_callable, deserialize_secrets_inp
 logger = logging.getLogger(__name__)
 
 
+@component
 class AzureOpenAIGenerator(OpenAIGenerator):
     """
     A Generator component that uses OpenAI's large language models (LLMs) on Azure to generate text.
diff --git a/haystack/components/generators/chat/azure.py b/haystack/components/generators/chat/azure.py
index 189b0c94..b6cd2e15 100644
--- a/haystack/components/generators/chat/azure.py
+++ b/haystack/components/generators/chat/azure.py
@@ -8,7 +8,7 @@ from typing import Any, Callable, Dict, Optional
 # pylint: disable=import-error
 from openai.lib.azure import AzureOpenAI
 
-from haystack import default_from_dict, default_to_dict, logging
+from haystack import component, default_from_dict, default_to_dict, logging
 from haystack.components.generators.chat import OpenAIChatGenerator
 from haystack.dataclasses import StreamingChunk
 from haystack.utils import Secret, deserialize_callable, deserialize_secrets_inplace, serialize_callable
@@ -16,6 +16,7 @@ from haystack.utils import Secret, deserialize_callable, deserialize_secrets_inp
 logger = logging.getLogger(__name__)
 
 
+@component
 class AzureOpenAIChatGenerator(OpenAIChatGenerator):
     """
     A Chat Generator component that uses the Azure OpenAI API to generate text.
diff --git a/haystack/core/pipeline/base.py b/haystack/core/pipeline/base.py
index a01423aa..e683f46c 100644
--- a/haystack/core/pipeline/base.py
+++ b/haystack/core/pipeline/base.py
@@ -5,6 +5,7 @@
 import importlib
 import itertools
 from collections import defaultdict
+from copy import copy, deepcopy
 from datetime import datetime
 from pathlib import Path
 from typing import Any, Dict, Iterator, List, Optional, TextIO, Tuple, Type, TypeVar, Union
@@ -640,45 +641,101 @@ class PipelineBase:
                         f"Input {socket_name} for component {component_name} is already sent by {socket.senders}."
                     )
 
-    def _prepare_component_input_data(self, data: Dict[str, Any]) -> Tuple[Dict[str, Dict[str, Any]], Dict[str, Any]]:
+    def _prepare_component_input_data(self, data: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
         """
         Prepares input data for pipeline components.
 
         Organizes input data for pipeline components and identifies any inputs that are not matched to any
-        component's input slots.
+        component's input slots. Deep-copies data items to avoid sharing mutables across multiple components.
 
         This method processes a flat dictionary of input data, where each key-value pair represents an input name
         and its corresponding value. It distributes these inputs to the appropriate pipeline components based on
         their input requirements. Inputs that don't match any component's input slots are classified as unresolved.
 
         :param data:
-            A dictionary with input names as keys and input values as values.
-        :returns:
-            A tuple containing two elements:
-             1. A dictionary mapping component names to their respective matched inputs.
-             2. A dictionary of inputs that were not matched to any component, termed as unresolved keyword arguments.
-        """
-        pipeline_input_data: Dict[str, Dict[str, Any]] = defaultdict(dict)
-        unresolved_kwargs = {}
+            A dictionary potentially having input names as keys and input values as values.
 
-        # Retrieve the input slots for each component in the pipeline
-        available_inputs: Dict[str, Dict[str, Any]] = self.inputs()
+        :returns:
+            A dictionary mapping component names to their respective matched inputs.
+        """
+        # check whether the data is a nested dictionary of component inputs where each key is a component name
+        # and each value is a dictionary of input parameters for that component
+        is_nested_component_input = all(isinstance(value, dict) for value in data.values())
+        if not is_nested_component_input:
+            # flat input, a dict where keys are input names and values are the corresponding values
+            # we need to convert it to a nested dictionary of component inputs and then run the pipeline
+            # just like in the previous case
+            pipeline_input_data: Dict[str, Dict[str, Any]] = defaultdict(dict)
+            unresolved_kwargs = {}
+
+            # Retrieve the input slots for each component in the pipeline
+            available_inputs: Dict[str, Dict[str, Any]] = self.inputs()
+
+            # Go through all provided to distribute them to the appropriate component inputs
+            for input_name, input_value in data.items():
+                resolved_at_least_once = False
+
+                # Check each component to see if it has a slot for the current kwarg
+                for component_name, component_inputs in available_inputs.items():
+                    if input_name in component_inputs:
+                        # If a match is found, add the kwarg to the component's input data
+                        pipeline_input_data[component_name][input_name] = input_value
+                        resolved_at_least_once = True
+
+                if not resolved_at_least_once:
+                    unresolved_kwargs[input_name] = input_value
+
+            if unresolved_kwargs:
+                logger.warning(
+                    "Inputs {input_keys} were not matched to any component inputs, please check your run parameters.",
+                    input_keys=list(unresolved_kwargs.keys()),
+                )
 
-        # Go through all provided to distribute them to the appropriate component inputs
-        for input_name, input_value in data.items():
-            resolved_at_least_once = False
+            data = dict(pipeline_input_data)
 
-            # Check each component to see if it has a slot for the current kwarg
-            for component_name, component_inputs in available_inputs.items():
-                if input_name in component_inputs:
-                    # If a match is found, add the kwarg to the component's input data
-                    pipeline_input_data[component_name][input_name] = input_value
-                    resolved_at_least_once = True
+        # deepcopying the inputs prevents the Pipeline run logic from being altered unexpectedly
+        # when the same input reference is passed to multiple components.
+        for component_name, component_inputs in data.items():
+            data[component_name] = {k: deepcopy(v) for k, v in component_inputs.items()}
 
-            if not resolved_at_least_once:
-                unresolved_kwargs[input_name] = input_value
+        return data
 
-        return pipeline_input_data, unresolved_kwargs
+    def _init_inputs_state(self, data: Dict[str, Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
+        for component_name, component_inputs in data.items():
+            if component_name not in self.graph.nodes:
+                # This is not a component name, it must be the name of one or more input sockets.
+                # Those are handled in a different way, so we skip them here.
+                continue
+            instance = self.graph.nodes[component_name]["instance"]
+            for component_input, input_value in component_inputs.items():
+                # Handle mutable input data
+                data[component_name][component_input] = copy(input_value)
+                if instance.__haystack_input__._sockets_dict[component_input].is_variadic:
+                    # Components that have variadic inputs need to receive lists as input.
+                    # We don't want to force the user to always pass lists, so we convert single values to lists here.
+                    # If it's already a list we assume the component takes a variadic input of lists, so we
+                    # convert it in any case.
+                    data[component_name][component_input] = [input_value]
+
+        return {**data}
+
+    def _init_to_run(self) -> List[Tuple[str, Component]]:
+        to_run: List[Tuple[str, Component]] = []
+        for node_name in self.graph.nodes:
+            component = self.graph.nodes[node_name]["instance"]
+
+            if len(component.__haystack_input__._sockets_dict) == 0:
+                # Component has no input, can run right away
+                to_run.append((node_name, component))
+                continue
+
+            for socket in component.__haystack_input__._sockets_dict.values():
+                if not socket.senders or socket.is_variadic:
+                    # Component has at least one input not connected or is variadic, can run right away.
+                    to_run.append((node_name, component))
+                    break
+
+        return to_run
 
     @classmethod
     def from_template(
@@ -707,6 +764,11 @@ class PipelineBase:
             msg += f"Source:\n{rendered}"
             raise PipelineUnmarshalError(msg)
 
+    def _init_graph(self):
+        """Resets the visits count for each component"""
+        for node in self.graph.nodes:
+            self.graph.nodes[node]["visits"] = 0
+
 
 def _connections_status(
     sender_node: str, receiver_node: str, sender_sockets: List[OutputSocket], receiver_sockets: List[InputSocket]
diff --git a/haystack/core/pipeline/pipeline.py b/haystack/core/pipeline/pipeline.py
index 2224f8dc..ed60e6f0 100644
--- a/haystack/core/pipeline/pipeline.py
+++ b/haystack/core/pipeline/pipeline.py
@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from copy import copy, deepcopy
+from copy import deepcopy
 from typing import Any, Dict, List, Mapping, Optional, Set, Tuple
 
 from haystack import logging, tracing
@@ -31,7 +31,19 @@ class Pipeline(PipelineBase):
 
         :param data:
             A dictionary of inputs for the pipeline's components. Each key is a component name
-            and its value is a dictionary of that component's input parameters.
+            and its value is a dictionary of that component's input parameters:
+            ```
+            data = {
+                "comp1": {"input1": 1, "input2": 2},
+            }
+            ```
+            For convenience, this format is also supported when input names are unique:
+            ```
+            data = {
+                "input1": 1, "input2": 2,
+            }
+            ```
+
         :param debug:
             Set to True to collect and return debug information.
         :param include_outputs_from:
@@ -83,85 +95,26 @@ class Pipeline(PipelineBase):
         {'hello2': {'output': 'Hello, Hello, world!!'}}.
         """
         pipeline_running(self)
-        # NOTE: We're assuming data is formatted like so as of now
-        # data = {
-        #     "comp1": {"input1": 1, "input2": 2},
-        # }
-        #
-        # TODO: Support also this format:
-        # data = {
-        #     "input1": 1, "input2": 2,
-        # }
 
         # Reset the visits count for each component
-        for node in self.graph.nodes:
-            self.graph.nodes[node]["visits"] = 0
+        self._init_graph()
 
         # TODO: Remove this warmup once we can check reliably whether a component has been warmed up or not
         # As of now it's here to make sure we don't have failing tests that assume warm_up() is called in run()
         self.warm_up()
 
-        # check whether the data is a nested dictionary of component inputs where each key is a component name
-        # and each value is a dictionary of input parameters for that component
-        is_nested_component_input = all(isinstance(value, dict) for value in data.values())
-        if not is_nested_component_input:
-            # flat input, a dict where keys are input names and values are the corresponding values
-            # we need to convert it to a nested dictionary of component inputs and then run the pipeline
-            # just like in the previous case
-            data, unresolved_inputs = self._prepare_component_input_data(data)
-            if unresolved_inputs:
-                logger.warning(
-                    "Inputs {input_keys} were not matched to any component inputs, please check your run parameters.",
-                    input_keys=list(unresolved_inputs.keys()),
-                )
+        # normalize `data`
+        data = self._prepare_component_input_data(data)
 
         # Raise if input is malformed in some way
         self._validate_input(data)
-        # NOTE: The above NOTE and TODO are technically not true.
-        # This implementation of run supports only the first format, but the second format is actually
-        # never received by this method. It's handled by the `run()` method of the `Pipeline` class
-        # defined in `haystack/pipeline.py`.
-        # As of now we're ok with this, but we'll need to merge those two classes at some point.
-
-        # deepcopying the inputs prevents the Pipeline run logic from being altered unexpectedly
-        # when the same input reference is passed to multiple components.
-        for component_name, component_inputs in data.items():
-            data[component_name] = {k: deepcopy(v) for k, v in component_inputs.items()}
-
-        for component_name, component_inputs in data.items():
-            if component_name not in self.graph.nodes:
-                # This is not a component name, it must be the name of one or more input sockets.
-                # Those are handled in a different way, so we skip them here.
-                continue
-            instance = self.graph.nodes[component_name]["instance"]
-            for component_input, input_value in component_inputs.items():
-                # Handle mutable input data
-                data[component_name][component_input] = copy(input_value)
-                if instance.__haystack_input__._sockets_dict[component_input].is_variadic:
-                    # Components that have variadic inputs need to receive lists as input.
-                    # We don't want to force the user to always pass lists, so we convert single values to lists here.
-                    # If it's already a list we assume the component takes a variadic input of lists, so we
-                    # convert it in any case.
-                    data[component_name][component_input] = [input_value]
-
-        last_inputs: Dict[str, Dict[str, Any]] = {**data}
+
+        # Initialize the inputs state
+        last_inputs: Dict[str, Dict[str, Any]] = self._init_inputs_state(data)
 
         # Take all components that have at least 1 input not connected or is variadic,
         # and all components that have no inputs at all
-        to_run: List[Tuple[str, Component]] = []
-        for node_name in self.graph.nodes:
-            component = self.graph.nodes[node_name]["instance"]
-
-            if len(component.__haystack_input__._sockets_dict) == 0:
-                # Component has no input, can run right away
-                to_run.append((node_name, component))
-                continue
-
-            for socket in component.__haystack_input__._sockets_dict.values():
-                if not socket.senders or socket.is_variadic:
-                    # Component has at least one input not connected or is variadic, can run right away.
-                    to_run.append((node_name, component))
-                    break
+        to_run: List[Tuple[str, Component]] = self._init_to_run()
 
         # These variables are used to detect when we're stuck in a loop.
         # Stuck loops can happen when one or more components are waiting for input but
diff --git a/releasenotes/notes/fix-azure-generators-serialization-18fcdc9cbcb3732e.yaml b/releasenotes/notes/fix-azure-generators-serialization-18fcdc9cbcb3732e.yaml
new file mode 100644
index 00000000..071969e6
--- /dev/null
+++ b/releasenotes/notes/fix-azure-generators-serialization-18fcdc9cbcb3732e.yaml
@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    Azure generators components fixed, they were missing the `@component` decorator.


Test patch:
diff --git a/test/components/generators/chat/test_azure.py b/test/components/generators/chat/test_azure.py
index 4b92fdee..c9693caa 100644
--- a/test/components/generators/chat/test_azure.py
+++ b/test/components/generators/chat/test_azure.py
@@ -6,6 +6,7 @@ import os
 import pytest
 from openai import OpenAIError
 
+from haystack import Pipeline
 from haystack.components.generators.chat import AzureOpenAIChatGenerator
 from haystack.components.generators.utils import print_streaming_chunk
 from haystack.dataclasses import ChatMessage
@@ -80,6 +81,15 @@ class TestOpenAIChatGenerator:
             },
         }
 
+    def test_pipeline_serialization_deserialization(self, tmp_path, monkeypatch):
+        monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
+        generator = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
+        p = Pipeline()
+        p.add_component(instance=generator, name="generator")
+        p_str = p.dumps()
+        q = Pipeline.loads(p_str)
+        assert p.to_dict() == q.to_dict(), "Pipeline serialization/deserialization w/ AzureOpenAIChatGenerator failed."
+
     @pytest.mark.integration
     @pytest.mark.skipif(
         not os.environ.get("AZURE_OPENAI_API_KEY", None) and not os.environ.get("AZURE_OPENAI_ENDPOINT", None),
diff --git a/test/components/generators/test_azure.py b/test/components/generators/test_azure.py
index b2373eb6..d5d52b6f 100644
--- a/test/components/generators/test_azure.py
+++ b/test/components/generators/test_azure.py
@@ -2,6 +2,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 import os
+
+from haystack import Pipeline
 from haystack.utils.auth import Secret
 
 import pytest
@@ -83,6 +85,15 @@ class TestAzureOpenAIGenerator:
             },
         }
 
+    def test_pipeline_serialization_deserialization(self, tmp_path, monkeypatch):
+        monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
+        generator = AzureOpenAIGenerator(azure_endpoint="some-non-existing-endpoint")
+        p = Pipeline()
+        p.add_component(instance=generator, name="generator")
+        p_str = p.dumps()
+        q = Pipeline.loads(p_str)
+        assert p.to_dict() == q.to_dict(), "Pipeline serialization/deserialization with AzureOpenAIGenerator failed."
+
     @pytest.mark.integration
     @pytest.mark.skipif(
         not os.environ.get("AZURE_OPENAI_API_KEY", None) and not os.environ.get("AZURE_OPENAI_ENDPOINT", None),
diff --git a/test/core/pipeline/test_pipeline.py b/test/core/pipeline/test_pipeline.py
index f91605cb..5c619d9f 100644
--- a/test/core/pipeline/test_pipeline.py
+++ b/test/core/pipeline/test_pipeline.py
@@ -14,7 +14,7 @@ from haystack.components.others import Multiplexer
 from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
 from haystack.components.routers import ConditionalRouter
 from haystack.core.component import component
-from haystack.core.component.types import InputSocket, OutputSocket
+from haystack.core.component.types import InputSocket, OutputSocket, Variadic
 from haystack.core.errors import PipelineDrawingError, PipelineError, PipelineMaxLoops, PipelineRuntimeError
 from haystack.core.pipeline import Pipeline, PredefinedPipeline
 from haystack.core.serialization import DeserializationCallbacks
@@ -175,7 +175,7 @@ def test_pipeline_resolution_duplicate_input_names_across_components():
     result = pipe.run(data={"what": "Haystack", "who": "world"})
     assert result == {"hello2": {"output": "Hello Hello world Haystack! Haystack!"}}
 
-    resolved, _ = pipe._prepare_component_input_data(data={"what": "Haystack", "who": "world"})
+    resolved = pipe._prepare_component_input_data(data={"what": "Haystack", "who": "world"})
 
     # why does hello2 have only one input? Because who of hello2 is inserted from hello.output
     assert resolved == {"hello": {"what": "Haystack", "who": "world"}, "hello2": {"what": "Haystack"}}
@@ -1008,3 +1008,82 @@ def test_pipeline_is_not_stuck_with_components_with_only_defaults_as_first_compo
     res = pipe.run({"prompt_builder": {"query": "What is the capital of Italy?"}})
 
     assert res == {"router": {"correct_replies": ["Rome"]}}
+
+
+def test__init_graph():
+    pipe = Pipeline()
+    pipe.add_component("greet", Greet())
+    pipe.add_component("adder", AddFixedValue())
+    pipe.connect("greet", "adder")
+    pipe._init_graph()
+    for node in pipe.graph.nodes:
+        assert pipe.graph.nodes[node]["visits"] == 0
+
+
+def test__init_to_run():
+    ComponentWithVariadic = component_class(
+        "ComponentWithVariadic", input_types={"in": Variadic[int]}, output_types={"out": int}
+    )
+    ComponentWithNoInputs = component_class("ComponentWithNoInputs", input_types={}, output_types={"out": int})
+    ComponentWithSingleInput = component_class(
+        "ComponentWithSingleInput", input_types={"in": int}, output_types={"out": int}
+    )
+    ComponentWithMultipleInputs = component_class(
+        "ComponentWithMultipleInputs", input_types={"in1": int, "in2": int}, output_types={"out": int}
+    )
+
+    pipe = Pipeline()
+    pipe.add_component("with_variadic", ComponentWithVariadic())
+    pipe.add_component("with_no_inputs", ComponentWithNoInputs())
+    pipe.add_component("with_single_input", ComponentWithSingleInput())
+    pipe.add_component("another_with_single_input", ComponentWithSingleInput())
+    pipe.add_component("with_multiple_inputs", ComponentWithMultipleInputs())
+
+    pipe.connect("with_no_inputs.out", "with_variadic.in")
+    pipe.connect("with_single_input.out", "another_with_single_input.in")
+    pipe.connect("another_with_single_input.out", "with_multiple_inputs.in1")
+    pipe.connect("with_multiple_inputs.out", "with_variadic.in")
+
+    to_run = pipe._init_to_run()
+    assert len(to_run) == 4
+    assert to_run[0][0] == "with_variadic"
+    assert to_run[1][0] == "with_no_inputs"
+    assert to_run[2][0] == "with_single_input"
+    assert to_run[3][0] == "with_multiple_inputs"
+
+
+def test__init_inputs_state():
+    pipe = Pipeline()
+    template = """
+    Answer the following questions:
+    {{ questions | join("\n") }}
+    """
+    pipe.add_component("prompt_builder", PromptBuilder(template=template))
+    pipe.add_component("multiplexer", Multiplexer(type_=int))
+    questions = ["What is the capital of Italy?", "What is the capital of France?"]
+    data = {
+        "prompt_builder": {"questions": questions},
+        "multiplexer": {"value": 1},
+        "not_a_component": "some input data",
+    }
+    res = pipe._init_inputs_state(data)
+    assert res == {
+        "prompt_builder": {"questions": ["What is the capital of Italy?", "What is the capital of France?"]},
+        "multiplexer": {"value": [1]},
+        "not_a_component": "some input data",
+    }
+    assert id(questions) != id(res["prompt_builder"]["questions"])
+
+
+def test__prepare_component_input_data():
+    MockComponent = component_class("MockComponent", input_types={"x": List[str], "y": str})
+    pipe = Pipeline()
+    pipe.add_component("first_mock", MockComponent())
+    pipe.add_component("second_mock", MockComponent())
+
+    res = pipe._prepare_component_input_data({"x": ["some data"], "y": "some other data"})
+    assert res == {
+        "first_mock": {"x": ["some data"], "y": "some other data"},
+        "second_mock": {"x": ["some data"], "y": "some other data"},
+    }
+    assert id(res["first_mock"]["x"]) != id(res["second_mock"]["x"])