我用这个 AI 工具生成单元测试，简直不要太爽！

本文分享如何使用驭码CodeRider 的单元测试功能生成单元测试文件。

在之前的文章如何用 Python 手撸一个 GitLab 代码安全审查工具？中，我用 Python 写了一个接受极狐GitLab 代码安全审计事件流并且将消息推送到钉钉群的脚本，完整的 python 代码为：

    from fastapi import FastAPIimport uvicornimport requestsimport jsonapp = FastAPI()@app.post("/jh-gitlab")async def gitlab_payload(data: dict):# 抓取审计事件中的主要信息audit_event_info = {"Action": data['details']['custom_message'],"Author": data['details']['author_name'],"IP Address": data['details']['ip_address'],"Entity Path": data['details']['entity_path'],"Target Details": data['target_details']}print(audit_event_info)# 发送消息通知notification(audit_event_info)def notification(payload: dict):webhook_url = "https://oapi.dingtalk.com/robot/send?access_token=你的钉钉 webhook token"# 发送消息的内容message = {"msgtype": "text","text": {"content" : "GitLab: {}".format(json.dumps(payload))}}# 发送 POST 请求headers = {'Content-Type': 'application/json'}response = requests.post(webhook_url, data=json.dumps(message), headers=headers)print(response.text)if json.loads(response.text)['errcode'] == 0:print("Send Message Success!")return Trueelse:print("Send Message Failed!")return json.loads(response.text)['errmsg']if __name__ == "__main__":uvicorn.run(app, host="0.0.0.0", port=8000)

对于这段代码写的是否 OK 最好的验证方法就是生成一些单元测试来测试。为了方便，我决定用 AI 工具来帮我生成单元测试代码。

我选择的是极狐GitLab 公司自研发布的 AI 产品驭码CodeRider（当前该产品可以免费试用【https://dl.gitlab.cn/w84csrf4】）。 该产品就有一个单元测试生成功能，当我在 PyCharm 中安装好驭码CodeRier 并直接点击test命令时：

就开始了单元测试代码的自动生成：

最终生成了如下测试代码：

from fastapi.testclient import TestClient
from unittest.mock import patch
import jsonfrom main import app, notificationclient = TestClient(app)def test_gitlab_payload_success():payload = {"details": {"custom_message": "Project marked for deletion","author_name": "Administrator","ip_address": "36.133.246.166","entity_path": "devsecops/ai-deleted-7"},"target_details": "ai-deleted-7"}with patch('main.notification') as mock_notification:mock_notification.return_value = Trueresponse = client.post("/jh-gitlab", json=payload)assert response.status_code == 200mock_notification.assert_called_once_with({"Action": "Project marked for deletion","Author": "Administrator","IP Address": "36.133.246.166","Entity Path": "devsecops/ai-deleted-7","Target Details": "ai-deleted-7"})def test_gitlab_payload_missing_fields():payload = {"details": {"custom_message": "Project marked for deletion","author_name": "Administrator","ip_address": "36.133.246.166"},"target_details": "ai-deleted-7"}response = client.post("/jh-gitlab", json=payload)assert response.status_code == 422def test_notification_success():payload = {"Action": "Project marked for deletion","Author": "Administrator","IP Address": "36.133.246.166","Entity Path": "devsecops/ai-deleted-7","Target Details": "ai-deleted-7"}with patch('requests.post') as mock_post:mock_response = mock_post.return_valuemock_response.text = json.dumps({"errcode": 0, "errmsg": "ok"})result = notification(payload)assert result == Truemock_post.assert_called_once()def test_notification_failure():payload = {"Action": "Project marked for deletion","Author": "Administrator","IP Address": "36.133.246.166","Entity Path": "devsecops/ai-deleted-7","Target Details": "ai-deleted-7"}with patch('requests.post') as mock_post:mock_response = mock_post.return_valuemock_response.text = json.dumps({"errcode": 400, "errmsg": "Bad Request"})result = notification(payload)assert result == "Bad Request"mock_post.assert_called_once()

看到这些代码，我就想知道怎么测试呢。然后就问驭码CodeRider，如果要运行这些单元测试，我该怎么操作？

驭码CodeRider 给的答案是四步：

• 第一步安装依赖
• 第二步将测试代码存放在 test_main.py文件中
• 第三步运行测试文件
• 第四步查看答案

我就遵照四步进行了测试，出现了错误：

提示 FAILED pytest_main.py::test_gitlab_payload_missing_entity_path - KeyError: 'entity_path'。仔细看了一下，测试代码中有一个检测缺失字段的环节，我代码中的 payload 有五个参数：Action、Author、IP Address、Entity Path以及 Target Details。

下面代码

def test_gitlab_payload_missing_fields():payload = {"details": {"custom_message": "Project marked for deletion","author_name": "Administrator","ip_address": "36.133.246.166"},"target_details": "ai-deleted-7"}response = client.post("/jh-gitlab", json=payload)assert response.status_code == 422

用来测试在缺失 entity_path字段的情况。比较遗憾的是，我在源代码中并没有对 payload 中的字段进行校验处理。所以我把这个错误发给了驭码CodeRider：

驭码CodeRider 给出了两种解决方案：

• 方案一：在测试中添加 entity_path 字段
• 方案二：修改 gitlab_payload 函数以处理缺失字段

按照这两种方式都可以，我选择了修改 gitlab_payload相关代码，于是继续问了驭码CodeRider：

驭码给的修改代码为：

    audit_event_info = {"Action": data['details'].get('custom_message', 'Unknown Action'),"Author": data['details'].get('author_name', 'Unknown Author'),"IP Address": data['details'].get('ip_address', 'Unknown IP'),"Entity Path": data['details'].get('entity_path', 'Unknown Path'),"Target Details": data.get('target_details', 'Unknown Target')}

就是给缺失的字段增加默认值。接着执行测试命令：

可以看到 4 条测试全部通过。

当然，上面的整个流程仅仅为测试使用，生成的单元测试不一定是最准确、最后直接可以使用的，但是我们可以看到用 AI 来生成单元测试文件至少是靠谱的、能够减轻不少工作量，先用 AI 生成，然后做一些修改，这样工作能轻松不少。

用 AI 来帮助生成单元测试文件看来靠谱，驭码CodeRider 可以的！

本文由博客群发一文多发等运营工具平台 OpenWrite 发布