Update test_stream_agent.py

test_stream_agent.py

@@ -1,18 +1,24 @@
-from stream_agent import StreamingAIAgent
+from langchain_core.messages import AIMessage
+
+from stream_agent import format_message
 
 
-def test_stream_returns_chunks(monkeypatch):
+def test_format_message_with_content():
-    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    message = AIMessage(content="hello")
-    agent = StreamingAIAgent()
 
-    chunks = list(agent.stream("test question"))
+    assert format_message(message) == "hello"
-
-    assert len(chunks) > 3
-    assert "".join(chunks).strip()
 
 
-def test_invoke_matches_stream(monkeypatch):
+def test_format_message_with_tool_call():
-    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    message = AIMessage(
-    agent = StreamingAIAgent()
+        content="",
+        tool_calls=[
+            {
+                "name": "get_demo_price",
+                "args": {"product": "молоко", "city": "Казань"},
+                "id": "call-1",
+            }
+        ],
+    )
 
-    assert agent.invoke("hello") == "".join(agent.stream("hello"))
+    assert "get_demo_price" in format_message(message)