我用這個 AI 工具生成單元測試，簡直不要太爽！

极狐GitLab發表於2024-10-30
AI

本文分享如何使用馭碼CodeRider 的單元測試功能生成單元測試檔案。

馭碼CodeRider 是極狐GitLab 自研的 AI 產品，可以進行 AI 程式設計和 DevOps 流程處理以及知識問答等，可申請免費試用【https://dl.gitlab.cn/bw72mqtc】

在之前的文章如何用 Python 手擼一個 GitLab 程式碼安全審查工具？中，我用 Python 寫了一個接受極狐GitLab 程式碼安全審計事件流並且將訊息推送到釘釘群的指令碼，完整的 python 程式碼為：

    from fastapi import FastAPI
    import uvicorn
    import requests
    import json

    app = 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 True
        else:
            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/bw72mqtc】）。 該產品就有一個單元測試生成功能，當我在 PyCharm 中安裝好馭碼CodeRier 並直接點選test命令時：

file

就開始了單元測試程式碼的自動生成：

file

最終生成了如下測試程式碼：

from fastapi.testclient import TestClient
from unittest.mock import patch
import json

from main import app, notification

client = 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 = True
        response = client.post("/jh-gitlab", json=payload)
        assert response.status_code == 200
        mock_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 == 422

def 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_value
        mock_response.text = json.dumps({"errcode": 0, "errmsg": "ok"})
        result = notification(payload)
        assert result == True
        mock_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_value
        mock_response.text = json.dumps({"errcode": 400, "errmsg": "Bad Request"})
        result = notification(payload)
        assert result == "Bad Request"
        mock_post.assert_called_once()

看到這些程式碼，我就想知道怎麼測試呢。然後就問馭碼CodeRider，如果要執行這些單元測試，我該怎麼操作？

file

馭碼CodeRider 給的答案是四步：

  • 第一步安裝依賴
  • 第二步將測試程式碼存放在 test_main.py檔案中
  • 第三步執行測試檔案
  • 第四步檢視答案

我就遵照四步進行了測試，出現了錯誤：

file

提示 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：

file

馭碼CodeRider 給出了兩種解決方案：

  • 方案一：在測試中新增 entity_path 欄位
  • 方案二：修改 gitlab_payload 函式以處理缺失欄位

按照這兩種方式都可以，我選擇了修改 gitlab_payload相關程式碼，於是繼續問了馭碼CodeRider：

file

馭碼給的修改程式碼為：

    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')
    }

就是給缺失的欄位增加預設值。接著執行測試命令：

file

可以看到 4 條測試全部透過。

當然，上面的整個流程僅僅為測試使用，生成的單元測試不一定是最準確、最後直接可以使用的，但是我們可以看到用 AI 來生成單元測試檔案至少是靠譜的、能夠減輕不少工作量，先用 AI 生成，然後做一些修改，這樣工作能輕鬆不少。

用 AI 來幫助生成單元測試檔案看來靠譜，馭碼CodeRider 可以的！

相關文章