Merge pull request #663 from iorisa/feature/unittest

feat: +unit test

tests/metagpt/learn/test_text_to_image.py

@@ -15,20 +15,24 @@ from metagpt.learn.text_to_image import text_to_image
 
 
 @pytest.mark.asyncio
-async def test():
+async def test_metagpt_llm():
     # Prerequisites
     assert CONFIG.METAGPT_TEXT_TO_IMAGE_MODEL_URL
     assert CONFIG.OPENAI_API_KEY
 
     data = await text_to_image("Panda emoji", size_type="512x512")
     assert "base64" in data or "http" in data
-    key = CONFIG.METAGPT_TEXT_TO_IMAGE_MODEL_URL
-    CONFIG.METAGPT_TEXT_TO_IMAGE_MODEL_URL = None
+
+    # Mock session env
+    old_options = CONFIG.options.copy()
+    new_options = old_options.copy()
+    new_options["METAGPT_TEXT_TO_IMAGE_MODEL_URL"] = None
+    CONFIG.set_context(new_options)
     try:
         data = await text_to_image("Panda emoji", size_type="512x512")
         assert "base64" in data or "http" in data
     finally:
-        CONFIG.METAGPT_TEXT_TO_IMAGE_MODEL_URL = key
+        CONFIG.set_context(old_options)
 
 
 if __name__ == "__main__":

tests/metagpt/learn/test_text_to_speech.py

@@ -27,13 +27,16 @@ async def test_text_to_speech():
     assert "base64" in data or "http" in data
 
     # test iflytek
-    key = CONFIG.AZURE_TTS_SUBSCRIPTION_KEY
-    CONFIG.AZURE_TTS_SUBSCRIPTION_KEY = ""
+    ## Mock session env
+    old_options = CONFIG.options.copy()
+    new_options = old_options.copy()
+    new_options["AZURE_TTS_SUBSCRIPTION_KEY"] = ""
+    CONFIG.set_context(new_options)
     try:
         data = await text_to_speech("panda emoji")
         assert "base64" in data or "http" in data
     finally:
-        CONFIG.AZURE_TTS_SUBSCRIPTION_KEY = key
+        CONFIG.set_context(old_options)
 
 
 if __name__ == "__main__":

tests/metagpt/roles/test_architect.py

@@ -7,18 +7,38 @@
 @Modified By: mashenquan, 2023-11-1. In accordance with Chapter 2.2.1 and 2.2.2 of RFC 116, utilize the new message
         distribution feature for message handling.
 """
+import uuid
+
 import pytest
 
+from metagpt.actions import WriteDesign, WritePRD
+from metagpt.config import CONFIG
+from metagpt.const import PRDS_FILE_REPO
 from metagpt.logs import logger
 from metagpt.roles import Architect
+from metagpt.schema import Message
+from metagpt.utils.common import any_to_str, awrite
 from tests.metagpt.roles.mock import MockMessages
 
 
 @pytest.mark.asyncio
 async def test_architect():
-    # FIXME: make git as env? Or should we support
+    # Prerequisites
+    filename = uuid.uuid4().hex + ".json"
+    await awrite(CONFIG.git_repo.workdir / PRDS_FILE_REPO / filename, data=MockMessages.prd.content)
+
     role = Architect()
-    role.put_message(MockMessages.req)
-    rsp = await role.run(MockMessages.prd)
+    rsp = await role.run(with_message=Message(content="", cause_by=WritePRD))
     logger.info(rsp)
     assert len(rsp.content) > 0
+    assert rsp.cause_by == any_to_str(WriteDesign)
+
+    # test update
+    rsp = await role.run(with_message=Message(content="", cause_by=WritePRD))
+    assert rsp
+    assert rsp.cause_by == any_to_str(WriteDesign)
+    assert len(rsp.content) > 0
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-s"])

tests/metagpt/roles/test_qa_engineer.py

@@ -5,3 +5,59 @@
 @Author  : alexanderwu
 @File    : test_qa_engineer.py
 """
+from pathlib import Path
+from typing import List
+
+import pytest
+from pydantic import Field
+
+from metagpt.actions import DebugError, RunCode, WriteTest
+from metagpt.actions.summarize_code import SummarizeCode
+from metagpt.config import CONFIG
+from metagpt.environment import Environment
+from metagpt.roles import QaEngineer
+from metagpt.schema import Message
+from metagpt.utils.common import any_to_str, aread, awrite
+
+
+async def test_qa():
+    # Prerequisites
+    demo_path = Path(__file__).parent / "../../data/demo_project"
+    CONFIG.src_workspace = Path(CONFIG.git_repo.workdir) / "qa/game_2048"
+    data = await aread(filename=demo_path / "game.py", encoding="utf-8")
+    await awrite(filename=CONFIG.src_workspace / "game.py", data=data, encoding="utf-8")
+    await awrite(filename=Path(CONFIG.git_repo.workdir) / "requirements.txt", data="")
+
+    class MockEnv(Environment):
+        msgs: List[Message] = Field(default_factory=list)
+
+        def publish_message(self, message: Message, peekable: bool = True) -> bool:
+            self.msgs.append(message)
+            return True
+
+    env = MockEnv()
+
+    role = QaEngineer()
+    role.set_env(env)
+    await role.run(with_message=Message(content="", cause_by=SummarizeCode))
+    assert env.msgs
+    assert env.msgs[0].cause_by == any_to_str(WriteTest)
+    msg = env.msgs[0]
+    env.msgs.clear()
+    await role.run(with_message=msg)
+    assert env.msgs
+    assert env.msgs[0].cause_by == any_to_str(RunCode)
+    msg = env.msgs[0]
+    env.msgs.clear()
+    await role.run(with_message=msg)
+    assert env.msgs
+    assert env.msgs[0].cause_by == any_to_str(DebugError)
+    msg = env.msgs[0]
+    env.msgs.clear()
+    role.test_round_allowed = 1
+    rsp = await role.run(with_message=msg)
+    assert "Exceeding" in rsp.content
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-s"])

tests/metagpt/test_role.py

@@ -10,15 +10,17 @@
     functionality is to be consolidated into the `Environment` class.
 """
 import uuid
+from unittest.mock import MagicMock
 
 import pytest
 from pydantic import BaseModel
 
 from metagpt.actions import Action, ActionOutput, UserRequirement
 from metagpt.environment import Environment
+from metagpt.provider.base_llm import BaseLLM
 from metagpt.roles import Role
 from metagpt.schema import Message
-from metagpt.utils.common import any_to_str
+from metagpt.utils.common import any_to_name, any_to_str
 
 
 class MockAction(Action):
@@ -96,7 +98,7 @@ async def test_react():
 
 
 @pytest.mark.asyncio
-async def test_msg_to():
+async def test_send_to():
     m = Message(content="a", send_to=["a", MockRole, Message])
     assert m.send_to == {"a", any_to_str(MockRole), any_to_str(Message)}
 
@@ -107,5 +109,50 @@ async def test_msg_to():
     assert m.send_to == {"a", any_to_str(MockRole), any_to_str(Message)}
 
 
+def test_init_action():
+    role = Role()
+    role.init_actions([MockAction, MockAction])
+    assert role.action_count == 2
+
+
+@pytest.mark.asyncio
+async def test_recover():
+    # Mock LLM actions
+    mock_llm = MagicMock(spec=BaseLLM)
+    mock_llm.aask.side_effect = ["1"]
+
+    role = Role()
+    assert role.is_watch(any_to_str(UserRequirement))
+    role.put_message(None)
+    role.publish_message(None)
+
+    role.llm = mock_llm
+    role.init_actions([MockAction, MockAction])
+    role.recovered = True
+    role.latest_observed_msg = Message(content="recover_test")
+    role.rc.state = 0
+    assert role.todo == any_to_name(MockAction)
+
+    rsp = await role.run()
+    assert rsp.cause_by == any_to_str(MockAction)
+
+
+@pytest.mark.asyncio
+async def test_think_act():
+    # Mock LLM actions
+    mock_llm = MagicMock(spec=BaseLLM)
+    mock_llm.aask.side_effect = ["ok"]
+
+    role = Role()
+    role.init_actions([MockAction])
+    await role.think()
+    role.rc.memory.add(Message("run"))
+    assert len(role.get_memories()) == 1
+    rsp = await role.act()
+    assert rsp
+    assert isinstance(rsp, ActionOutput)
+    assert rsp.content == "run"
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-s"])

tests/metagpt/test_schema.py

@@ -16,8 +16,10 @@ from metagpt.actions import Action
 from metagpt.actions.action_node import ActionNode
 from metagpt.actions.write_code import WriteCode
 from metagpt.config import CONFIG
+from metagpt.const import SYSTEM_DESIGN_FILE_REPO, TASK_FILE_REPO
 from metagpt.schema import (
     AIMessage,
+    CodeSummarizeContext,
     Document,
     Message,
     MessageQueue,
@@ -61,6 +63,8 @@ def test_message():
     assert m.role == "b"
     assert m.send_to == {"c"}
     assert m.cause_by == "c"
+    m.sent_from = "e"
+    assert m.sent_from == "e"
 
     m.cause_by = "Message"
     assert m.cause_by == "Message"
@@ -121,6 +125,8 @@ def test_document():
 @pytest.mark.asyncio
 async def test_message_queue():
     mq = MessageQueue()
+    val = await mq.dump()
+    assert val == "[]"
     mq.push(Message(content="1"))
     mq.push(Message(content="2中文测试aaa"))
     msg = mq.pop()
@@ -132,5 +138,23 @@ async def test_message_queue():
     assert new_mq.pop_all() == mq.pop_all()
 
 
+@pytest.mark.parametrize(
+    ("file_list", "want"),
+    [
+        (
+            [f"{SYSTEM_DESIGN_FILE_REPO}/a.txt", f"{TASK_FILE_REPO}/b.txt"],
+            CodeSummarizeContext(
+                design_filename=f"{SYSTEM_DESIGN_FILE_REPO}/a.txt", task_filename=f"{TASK_FILE_REPO}/b.txt"
+            ),
+        )
+    ],
+)
+def test_CodeSummarizeContext(file_list, want):
+    ctx = CodeSummarizeContext.loads(file_list)
+    assert ctx == want
+    m = {ctx: ctx}
+    assert want in m
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-s"])

tests/metagpt/utils/test_redis.py

@@ -27,6 +27,19 @@ async def test_redis():
     assert await conn.get("test") == b"test"
     await conn.close()
 
+    # Mock session env
+    old_options = CONFIG.options.copy()
+    new_options = old_options.copy()
+    new_options["REDIS_HOST"] = "YOUR_REDIS_HOST"
+    CONFIG.set_context(new_options)
+    try:
+        conn = Redis()
+        await conn.set("test", "test", timeout_sec=0)
+        assert not await conn.get("test") == b"test"
+        await conn.close()
+    finally:
+        CONFIG.set_context(old_options)
+
 
 if __name__ == "__main__":
     pytest.main([__file__, "-s"])

tests/metagpt/utils/test_s3.py

@@ -41,17 +41,18 @@ async def test_s3():
     res = await conn.cache(data, ".bak", "script")
     assert "http" in res
 
-
-@pytest.mark.asyncio
-async def test_s3_no_error():
-    conn = S3()
-    key = conn.auth_config["aws_secret_access_key"]
-    conn.auth_config["aws_secret_access_key"] = ""
+    # Mock session env
+    old_options = CONFIG.options.copy()
+    new_options = old_options.copy()
+    new_options["S3_ACCESS_KEY"] = "YOUR_S3_ACCESS_KEY"
+    CONFIG.set_context(new_options)
     try:
+        conn = S3()
+        assert not conn.is_valid
         res = await conn.cache("ABC", ".bak", "script")
         assert not res
     finally:
-        conn.auth_config["aws_secret_access_key"] = key
+        CONFIG.set_context(old_options)
 
 
 if __name__ == "__main__":