From d9c5809ccd6acc74aea3307c9ef0f5069f530c61 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=8E=98=E6=9D=83=20=E9=A9=AC?= <mashenquan@fuzhi.ai>
Date: Tue, 2 Jan 2024 14:21:55 +0800
Subject: [PATCH] feat: +qa unit test

---
 tests/data/demo_project/game.py         | 92 +++++++++++++++++++++++++
 tests/metagpt/roles/test_qa_engineer.py | 56 +++++++++++++++
 2 files changed, 148 insertions(+)
 create mode 100644 tests/data/demo_project/game.py

diff --git a/tests/data/demo_project/game.py b/tests/data/demo_project/game.py
new file mode 100644
index 000000000..22e77b260
--- /dev/null
+++ b/tests/data/demo_project/game.py
@@ -0,0 +1,92 @@
+## game.py
+
+import random
+from typing import List, Tuple
+
+
+class Game:
+    def __init__(self):
+        self.grid: List[List[int]] = [[0 for _ in range(4)] for _ in range(4)]
+        self.score: int = 0
+        self.game_over: bool = False
+
+    def reset_game(self):
+        self.grid = [[0 for _ in range(4)] for _ in range(4)]
+        self.score = 0
+        self.game_over = False
+        self.add_new_tile()
+        self.add_new_tile()
+
+    def move(self, direction: str):
+        if direction == "up":
+            self._move_up()
+        elif direction == "down":
+            self._move_down()
+        elif direction == "left":
+            self._move_left()
+        elif direction == "right":
+            self._move_right()
+
+    def is_game_over(self) -> bool:
+        for i in range(4):
+            for j in range(4):
+                if self.grid[i][j] == 0:
+                    return False
+                if j < 3 and self.grid[i][j] == self.grid[i][j + 1]:
+                    return False
+                if i < 3 and self.grid[i][j] == self.grid[i + 1][j]:
+                    return False
+        return True
+
+    def get_empty_cells(self) -> List[Tuple[int, int]]:
+        empty_cells = []
+        for i in range(4):
+            for j in range(4):
+                if self.grid[i][j] == 0:
+                    empty_cells.append((i, j))
+        return empty_cells
+
+    def add_new_tile(self):
+        empty_cells = self.get_empty_cells()
+        if empty_cells:
+            x, y = random.choice(empty_cells)
+            self.grid[x][y] = 2 if random.random() < 0.9 else 4
+
+    def get_score(self) -> int:
+        return self.score
+
+    def _move_up(self):
+        for j in range(4):
+            for i in range(1, 4):
+                if self.grid[i][j] != 0:
+                    for k in range(i, 0, -1):
+                        if self.grid[k - 1][j] == 0:
+                            self.grid[k - 1][j] = self.grid[k][j]
+                            self.grid[k][j] = 0
+
+    def _move_down(self):
+        for j in range(4):
+            for i in range(2, -1, -1):
+                if self.grid[i][j] != 0:
+                    for k in range(i, 3):
+                        if self.grid[k + 1][j] == 0:
+                            self.grid[k + 1][j] = self.grid[k][j]
+                            self.grid[k][j] = 0
+
+    def _move_left(self):
+        for i in range(4):
+            for j in range(1, 4):
+                if self.grid[i][j] != 0:
+                    for k in range(j, 0, -1):
+                        if self.grid[i][k - 1] == 0:
+                            self.grid[i][k - 1] = self.grid[i][k]
+                            self.grid[i][k] = 0
+
+    def _move_right(self):
+        for i in range(4):
+            for j in range(2, -1, -1):
+                if self.grid[i][j] != 0:
+                    for k in range(j, 3):
+                        if self.grid[i][k + 1] == 0:
+                            self.grid[i][k + 1] = self.grid[i][k]
+                            self.grid[i][k] = 0
diff --git a/tests/metagpt/roles/test_qa_engineer.py b/tests/metagpt/roles/test_qa_engineer.py
index 8fd7c0373..784c26a06 100644
--- a/tests/metagpt/roles/test_qa_engineer.py
+++ b/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"])