From 7e70f4b60802e6c4b92f086117f5050bf52ca783 Mon Sep 17 00:00:00 2001
From: Xingzhi Zhang <37076709+elliotzh@users.noreply.github.com>
Date: Thu, 9 May 2024 11:11:20 +0800
Subject: [PATCH] fix: pass attribute aether for AetherBridgeComponent (#35534)

---
 sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/component.py | 2 ++
 sdk/ml/azure-ai-ml/tests/internal/unittests/test_component.py  | 1 +
 2 files changed, 3 insertions(+)

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/component.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/component.py
index 810579cbf409..1040a12526a0 100644
--- a/sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/component.py
+++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/component.py
@@ -103,6 +103,7 @@ def __init__(
         ae365exepool: Optional[Dict] = None,
         launcher: Optional[Dict] = None,
         datatransfer: Optional[Dict] = None,
+        aether: Optional[Dict] = None,
         **kwargs,
     ):
         _type, self._type_label = parse_name_label(type)
@@ -141,6 +142,7 @@ def __init__(
         self.ae365exepool = ae365exepool
         self.launcher = launcher
         self.datatransfer = datatransfer
+        self.aether = aether
 
     @classmethod
     def _build_io(cls, io_dict: Union[Dict, Input, Output], is_input: bool):
diff --git a/sdk/ml/azure-ai-ml/tests/internal/unittests/test_component.py b/sdk/ml/azure-ai-ml/tests/internal/unittests/test_component.py
index df8ae6720446..99da4145bcc7 100644
--- a/sdk/ml/azure-ai-ml/tests/internal/unittests/test_component.py
+++ b/sdk/ml/azure-ai-ml/tests/internal/unittests/test_component.py
@@ -989,6 +989,7 @@ def test_load_from_internal_aether_bridge_component(self):
                 "mock_param1": {"type": "AnyFile", "optional": False},
                 "mock_param2": {"type": "AnyFile", "optional": False},
             },
+            "aether": {"module_type": "ScrapingCloud", "ref_id": "mock_ref_id"},
             "outputs": {"job_info": {"type": "AnyFile"}},
             "type": "AetherBridgeComponent",
             "command": "mock.exe {inputs.mock_param1} {inputs.mock_param2} {outputs.job_info}",