diff --git a/tests/test_agent_server.py b/tests/test_agent_server.py
index 9fae302e..d41b3865 100644
--- a/tests/test_agent_server.py
+++ b/tests/test_agent_server.py
@@ -1,20 +1,15 @@
 from llama_index.core.llms import MockLLM
 from llama_index.core.agent import ReActAgent
 
-from agentfile.agent_server import FastAPIAgentServer
+from agentfile.services import AgentService
 from agentfile.message_queues.simple import SimpleMessageQueue
 
-from unittest.mock import patch, MagicMock
 
-
-@patch("agentfile.agent_server.fastapi.uuid")
-def test_init(mock_uuid: MagicMock) -> None:
-    mock_uuid.uuid4.return_value = "mock"
+def test_init() -> None:
     agent = ReActAgent.from_tools([], llm=MockLLM())
-    mq = SimpleMessageQueue()
-    server = FastAPIAgentServer(
+    server = AgentService(
         agent,
-        mq,
+        SimpleMessageQueue(),
         running=False,
         description="Test Agent Server",
         step_interval=0.5,
@@ -24,5 +19,3 @@ def test_init(mock_uuid: MagicMock) -> None:
     assert server.running is False
     assert server.description == "Test Agent Server"
     assert server.step_interval == 0.5
-    assert server.message_queue == mq
-    assert server.publisher_id == "FastAPIAgentServer-mock"