AIGC 应用系列之 —— 单元测试

序

OpenAI 的确引发了 AIGC 的热潮。作为工程师，我们成为了第一批受益者，无需担心被取代。AIGC 带来的效能提升是显而易见的。在接下来的这一系列文章中，@Arno 将不再讨论如何实现那些普通工程师难以触及的成果，而是聚焦于当前的工作，为工程师日常遇到的场景利用 AI 提升效率提供最佳实践。

说到编写单元测试，大家第一反应都是：「真香」，但是「业务压力大」，能够用于写单测的时间并不多，ROI 不高，因此在大部分项目中选择放弃。而 AI 的到来，可以大幅缩短针对特定业务代码或者基础代码编写的时间，从而让核心项目快速覆盖单元测试变得唾手可得 ~

笔者接下来将以自己的项目为例，结合：Jest / Typescript / Node.js 的技术栈为例，讲解如何写好 PROMPT 并应用在自己的日常工作中。

🦄 沿用部分 TDD 的思想并结合 AI 的能力编写代码

在开始之前，不得不提一下 TDD。大型项目的核心代码，在开始编码的时候就可以考虑 TDD 驱动着做，因为这样会让你编写的代码具备如下的便于 AI 理解（当然也便于人理解）的特性：

• 高度模块化、模块之间松耦合，最好是面向接口编程。（AI 能够非常好地将接口作为上下文进行推理）
• 关注纯函数（无副作用 SideEffects 函数）的编写，纯函数的实现能够让 AI 更容易理解，编写的单元测试也将更加可靠。
• 控制类的长度和复杂度，控制函数的长度和复杂度（比如加一个 Lint 规则，超过 1k 行就提醒，并反思设计和实现的合理性，并考虑做重构、拆解和优化）。
• ...

简而言之，代码的编写需要做到简单优雅，甚至可以去追求函数式编程的「艺术」（开个玩笑） ~ 这样的代码往往能够让 AI 写出一次性就能够通过的 TestCases，想想是不是觉得很激动！

🪄 编写单元测试的 Prompt

首先我们来解决一个简单的单元测试编写 Case，参看下面的代码，如何为之编写单元测试的 PROMPT 呢？

import { injectable } from 'inversify';

import { createServiceInjector } from './di';

export const IKeyStore = createServiceInjector(Symbol('IKeyStore'));

export interface IKeyStore {
  set(key: string, value: string): void;
  get(key: string): string | undefined;
  setKeys(key: Record<string, string>): void;
}

@injectable()
export class KeyStore implements IKeyStore {
  private _store: Map<string, string> = new Map();

  public set(key: string, value: string) {
    this._store.set(key, value);
  }

  public get(key: string) {
    return this._store.get(key);
  }

  public setKeys(keys: Record<string, string>) {
    Object.keys(keys).forEach((key) => {
      this._store.set(key, keys[key]);
    });
  }
}

我们可以根据社区中，包括 www.yuque.com/surfacew/da… 等一系列的 Prompt 教程的指导下，我个人尝试编写了一个针对上述技术栈的 Prompt 生成器。

这个 Prompt 结构包括了：

• 语言：我们尽量配合 OpenAI 使用英文去描述 Prompt

• • 一者省 tokens
  • 二者 OpenAI 更擅长英文

• 角色：我们将之定义为一个代码生成器
• 指令：对于指令区域我们要求它做如下的事情，强调如下的指令：

• • 使用 JEST 作为测试框架
  • 使用 Typescript 作为编码语言
  • 运行环境为 Node.js
  • 只需要生成代码，不需要返回解释
  • 其它的针对运行特定的代码的补充：比如，只考虑生成 XX 类或者函数的测试代码

• 输出：指明输出为 Typescript

将这个 Prompt 整理出来，我们看一下效果：

Your current role is: UnitTest code generator
##
- use jest test framework
- use typescript
- the test environment is Node.js
- only generate code with no explain
- write the`prepareChatMessages` and `getFirstAnswerMsg` and `setChatMessages` and `setModel` method's unit test cases
Here is your target code to generate tests:
```
{{your_code_here}}
```
##
``` typescript ``` 

试着运行一下，你会发现，它在 OpenAI Turbo 3.5 的环境下输出的代码可以直接使用，非常惊艳！

import 'reflect-metadata';
import { IKeyStore, KeyStore } from './key-store';

describe('KeyStore', () => {
  let keyStore: IKeyStore;

  beforeEach(() => {
    keyStore = new KeyStore();
  });

  describe('set', () => {
    it('should set a value to the given key', () => {
      const key = 'key';
      const value = 'value';
      keyStore.set(key, value);
      expect(keyStore.get(key)).toEqual(value);
    });
  });

  describe('get', () => {
    it('should return the value for the given key', () => {
      const key = 'key';
      const value = 'value';
      keyStore.set(key, value);
      expect(keyStore.get(key)).toEqual(value);
    });

    it('should return undefined for a non-existent key', () => {
      const key = 'non-existent-key';
      expect(keyStore.get(key)).toBeUndefined();
    });
  });

  describe('setKeys', () => {
    it('should set multiple key-value pairs to the store', () => {
      const keys = {
        key1: 'value1',
        key2: 'value2',
        key3: 'value3',
      };
      keyStore.setKeys(keys);
      Object.keys(keys).forEach((key) => {
        expect(keyStore.get(key)).toEqual(keys[key]);
      });
    });
  });
});

OK，那么你会说：就这？这不就是个简单的 getter / setter 么？那么复杂一点的情形的单测有需要如何去优化呢？别急，且待我徐徐道来。

📦 输入长上下文压缩

对于部分上下文很多的代码，特别是依赖了大量外部的 module 或者内容的代码，轻松就超过了 4k tokens 了。为此我们要如何解决这个问题呢？

对于日常工作而言，我主要给出几种可行的简单思路：

• 使用更高级的 API：也就是 GPT 4.x，支持 8k / 32k 的上下文传参，不缺钱的不要犹豫直接上
• 压缩代码：使用线上工具或者本地工具，比如：www.novel.tools/minify/type…，将你自己的代码进行 Uglify 压缩，让上下文得到缩减，去除不必要的空白符、注释等 AI 并不需要关注的东西，大量降低 tokens。以上面的 KeyStore 为例，压缩前 token 数量为 251， 压缩后 195，如果有注释等代码收益会更加明显。

• 分离文件实现：通过重构的过程，尽量将单体的函数、类拆解为粒度更小的实现，便于做任务拆分，也同时降低文件对于开发者本身的认知复杂度。
• 使用 Continue 语句：对于部分 AI 尤其是 ChatGPT 无法一次性吐出的情形，我们可以在 chat 后使用 continue 的消息，让 AI 继续完成吐出测试 Code。

🌟 提升单测 TC 生成质量

对于部分复杂度较高的项目，一个文件可能包含大量的 import 内容，这个时候我们最好考虑使用如下的策略来做优化：

• 使用接口来替换实现，让 AI 理解 Interface 声明，而非去关注外部实现。比如下面的代码中：

import { IModelSchema } from '../sdk-share/lib';
import { Prompt } from '../../prompts/prompt';
import { ITaskConfigClient, MLModelInvokeParamType } from '../../tasks/task';
import { ILogger } from '../../utils/logger';
import { BaseMLModel } from '../model.base';
import { invokeOpenAIRpcStream } from './chat.openai.stream';

我们依赖了较多的接口、类或者函数。为了让 AI 理解这些上下文同时不至于大量增加上下文的数量，我们可以将 Prompt 和 BaseMLModel 这类实现转变为接口，在代码引用的地方使用 as IPrompt、as IBaseMLModel 而不用讲具体实现拷贝进入上下文中。通过这样的手段既给 AI 提供了充裕的上下文，又可以减少 tokens 的数量。

---

• 如果是使用 OpenAI 的 API 做生成（土豪任性），建议略微调低 temperature 的值，到 0.8 左右，降低发散，提升准确性。

---

• 可以多次 Retry 让 AI 生成一些差异化的 Cases 再进行筛选。找到差集的部分，手动引用或者进行补充。

---

• 当然也可以在提示词中指出：「尽可能多地考虑到边界条件」生成覆盖率接近 100% 」 的指令。以便于 AI 生成更多更细致的测试用例。

总结

AI 编写单测，能够让平时一天的编写，一小时就搞定，大大加速了我们保障软件质量所需要付出的代价。这块玩法 ROI 极高，期待大家的实践和效果分享，一起找到最佳实践！