From da526f3ae036cc16e25e9e793707fa0474bb80da Mon Sep 17 00:00:00 2001
From: Eitan Yarmush <eitan.yarmush@solo.io>
Date: Wed, 5 Feb 2025 21:31:04 +0000
Subject: [PATCH] fixes

---
 .../autogen-agentchat/tests/test_assistant_agent.py    | 10 +++++++---
 python/packages/autogen-core/tests/test_memory.py      |  3 +--
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/python/packages/autogen-agentchat/tests/test_assistant_agent.py b/python/packages/autogen-agentchat/tests/test_assistant_agent.py
index 4e2dbdeed571..040f3e23e6ca 100644
--- a/python/packages/autogen-agentchat/tests/test_assistant_agent.py
+++ b/python/packages/autogen-agentchat/tests/test_assistant_agent.py
@@ -3,7 +3,6 @@
 import logging
 from typing import Any, AsyncGenerator, List
 
-from autogen_core._component_config import ComponentModel
 import pytest
 from autogen_agentchat import EVENT_LOGGER_NAME
 from autogen_agentchat.agents import AssistantAgent
@@ -19,7 +18,7 @@
     ToolCallRequestEvent,
     ToolCallSummaryMessage,
 )
-from autogen_core import FunctionCall, Image
+from autogen_core import ComponentModel, FunctionCall, Image
 from autogen_core.memory import ListMemory, Memory, MemoryContent, MemoryMimeType, MemoryQueryResult
 from autogen_core.model_context import BufferedChatCompletionContext
 from autogen_core.models import (
@@ -755,7 +754,12 @@ async def test_run_with_memory(monkeypatch: pytest.MonkeyPatch) -> None:
         "test_agent", model_client=OpenAIChatCompletionClient(model=model, api_key=""), memory=[memory2]
     )
 
-    result = await agent.run(task="test task")
+    # Test dump and load component with memory
+    agent_config: ComponentModel = agent.dump_component()
+    assert agent_config.provider == "autogen_agentchat.agents.AssistantAgent"
+    agent2 = AssistantAgent.load_component(agent_config)
+
+    result = await agent2.run(task="test task")
     assert len(result.messages) > 0
     memory_event = next((msg for msg in result.messages if isinstance(msg, MemoryQueryEvent)), None)
     assert memory_event is not None
diff --git a/python/packages/autogen-core/tests/test_memory.py b/python/packages/autogen-core/tests/test_memory.py
index d3f5410feba2..ce98aaffb97f 100644
--- a/python/packages/autogen-core/tests/test_memory.py
+++ b/python/packages/autogen-core/tests/test_memory.py
@@ -1,8 +1,7 @@
 from typing import Any
 
 import pytest
-from autogen_core import CancellationToken
-from autogen_core._component_config import ComponentModel
+from autogen_core import CancellationToken, ComponentModel
 from autogen_core.memory import (
     ListMemory,
     Memory,