🏭

テストデータ生成器を自動生成

2024/12/13に公開

この記事は株式会社ビットキー Advent Calendar 2024 13日目の記事です。
技術本部の maruware が担当します。

概要

自動テストのコード内で使うテストデータを生成するファクトリをPrismaやOpenAPI specおよびZodスキーマから自動生成した話

背景

自動テストのコードを書いているとテスト用のデータを構成することがよくあります。

データにはたくさんのフィールドがあり、単純に記述するとテストの結果に影響のある(注目している)フィールドがどれなのかがわかりづらくなります。

また、requiredなフィールドをどうにか埋めるのも面倒なポイントです。

そのため、テストデータを返すユーティリティ関数を用意することがよくありました。しかしながら、あらゆるエンティティ用に実装していくのは非常に面倒です。

ビットキーで提供している workhub というプロダクトでは多くの仕組みをスキーマファーストで実装しています。

そこで、スキーマからテストデータファクトリを自動生成することで楽したいと考えました。

元々の状態

イメージとして、以下のようなデータ生成関数を実装していました。

export class UserDataCreator {
  public static createUser(params: Partial<User>): User => {
    return {
      id: randomUUID(),
      name: "bitkey user",
			// ...
      ...p,
    }));
  };
}

また、上記を使ってさらにデータベースへの挿入まで行う以下のような関数も用意されていました。

export async prepareUser(params?: Partial<User>): Promise<User> => {
  const userData = UserDataCreator.createUser({...params});
  return await db.create(userData);
};

上記について以下のような課題感を持っていました。

  1. いちいち用意するのが面倒
  2. インターフェースが画一的になっていない(画一的にするために人間の意識に依存)

workhub では上記のようなユーティリティをRDB ( Prisma )とFirestore用に用意していました。

検討したこと

まず、Prisma用の機構を用意することを検討しました。

このような用途のためのライブラリとして有名どころだとprisma-fabbricaがあります。

https://github.com/Quramy/prisma-fabbrica

基本的にはいい感じですが、workhubにおいては以下のような課題がありました。

  1. string型のカラムをTypeScript上でだけ文字列リテラルのUnionとしている疑似的なenumやJSON型の内容について、Zodスキーマを定義して拡張するようにしていて、Prismaスキーマだけだと不完全
  2. たとえばemailというフィールド名ならemailとして生成したデータを自動で入れるなどルールを拡張したい

1は具体的な例としては以下のようなカスタムを行っています。

model Participant {
  /// @zod.enum("Member", "Visitor", "Customer")
  type String
  
  /// @zodjson("participant_content.yaml")
  content Json
}

1の解決として、zodスキーマをベースにした仕組みを検討したところ、以下の @anatine/zod-mock などいくつかのライブラリが見つかりました。

https://www.npmjs.com/package/@anatine/zod-mock

しかし、上記だと我々がやりたいことを満たす上で少し拡張性が足りない部分がありました。

たとえば、workhubではPrismaのデータベースとしてPostgreSQLを使用していますが、PostgreSQLではintegerが32ビットであり最大値を制限したいなどがありました。

( @anatine/zod-mock ではfakerが使われており、fakerのint はデフォルトで最大が Number.MAX_SAFE_INTEGER になっている)

そのため、似たようなライブラリを自作することとしました。

やったこと

@anatine/zod-mock から以下を拡張したライブラリを用意しました。

  1. fillOptional オプションでoptionalフィールドも埋められるように
  2. constraints オプションで基本の制約を入れられるように(Int32用)

以降、このライブラリを zod-fake と呼称します。

次に、zod-fakeを使ったテストデータ生成器(ファクトリー)を生成するPrisma generatorを実装しました。

ファクトリーのインターフェースはprisma-fabbricaをインスパイアした以下のような形にしました。


import type {KeyToFake, FakerFunction, ZodFakeOptions} from '@bitkey-service/zod-fake';
import {fake} from '@bitkey-service/zod-fake';

export const commonKeyToFake: KeyToFake = (keyName, faker) => {
  const fixedKeyMap: {[keyName: string]: FakerFunction} = {
    nameJa: fakerJA.person.fullName,
    firstNameJa: fakerJA.person.firstName,
    familyNameJa: fakerJA.person.lastName,
    nameEn: faker.person.fullName,
    firstNameEn: faker.person.firstName,
    familyNameEn: faker.person.lastName,
    title: fakerJA.lorem.sentence,
    prefecture: fakerJA.location.state,
    city: fakerJA.location.city,
    country: fakerJA.location.country,
  };
  if (fixedKeyMap[keyName]) {
    return fixedKeyMap[keyName];
  }

  if (/phoneNumber$/i.test(keyName)) {
    return fakerJA.phone.number;
  }
  if (/^addressLine/.test(keyName)) {
    return fakerJA.location.streetAddress;
  }
  if (/url$/i.test(keyName)) {
    return faker.internet.url;
  }
};

export function defineUserFactory(db: DBClient, opts?: ZodFakeOptions) {
  const keyToFake: KeyToFake = (keyName, faker) => {
    if (opts?.keyToFake) {
      const ret = opts.keyToFake(keyName, faker);
      if (ret) return ret;
    }

    if (keyName === 'id') return () => faker.string.alphanumeric(20);

    return commonKeyToFake(keyName, faker);
  };
  const paramsSchema = userSchema.partial();
  const build = (params?: z.infer<typeof paramsSchema>, buildOpts?: ZodFakeOptions) => {
    const data = fake(userSchema, {faker, constraints: {int: {max: 2147483647}}, ...opts, ...buildOpts, keyToFake});
    return {...data, ...params};
  };
  const create = async (params?: z.infer<typeof paramsSchema>, buildOpts?: ZodFakeOptions) => {
    const data = build(params, buildOpts);
    return db.user.create({data: data});
  };
  const createMany = async (paramsList: z.infer<typeof paramsSchema>[], buildOpts?: ZodFakeOptions) => {
    const dataList = paramsList.map(params => build(params, buildOpts)).map(data => data);
    return db.user.createManyAndReturn({data: dataList});
  };
  return {
    build,
    create,
    createMany,
  };
}

以下のようなことをデフォルトとすることでより楽にしています。

  • workhubでよくあるフィールド名でfakerの関数をデフォルトで決めうつ
  • PostgreSQLのIntの32ビット上限をconstraintsとしてデフォルトで入れる
  • idがPrismaスキーマ上でuuidならfakerをuuidに、そうではないstringならalphanumericに

使う際は以下のような形になります。

const userFactory = defineUserFactory(db)

const userParam = userFactory.build({nameJa: '佐藤 雅彦'})
const user = await userFactory.create({nameJa: '佐藤 雅彦'}) 

Prisma generatorは主に以下のような実装です。

import type {DMMF} from '@prisma/generator-helper';
import {camelCase} from 'change-case';

export function makeContent (models: readonly DMMF.Model[]) {
  const modelNames = models.map(d => d.name);
  const importContents = [
    `import type {DBClient} from '@/common/dbClient';`,
    `import {Prisma} from '@workhub/prisma-client';`,
    `import type {KeyToFake, FakerFunction, ZodFakeOptions} from '@bitkey-service/zod-fake';`,
    `import {fake} from '@bitkey-service/zod-fake';`,
    `import type {z} from 'zod';`,
    `import { faker, fakerJA } from '@faker-js/faker';`,
    ...modelNames.map(modelName => {
      const camelModelName = camelCase(modelName);
      return `import {${camelModelName}Schema} from '@/prisma/entities/${camelModelName}Entity';`;
    }),
  ];

  const commonKeyToFakeContent = `
  export const commonKeyToFake: KeyToFake = (keyName, faker) => {
    const fixedKeyMap: {[keyName: string]: FakerFunction} = {
      nameJa: fakerJA.person.fullName,
      firstNameJa: fakerJA.person.firstName,
      familyNameJa: fakerJA.person.lastName,
      nameEn: faker.person.fullName,
      firstNameEn: faker.person.firstName,
      familyNameEn: faker.person.lastName,
      title: fakerJA.lorem.sentence,
      prefecture: fakerJA.location.state,
      city: fakerJA.location.city,
      country: fakerJA.location.country,
    }
    if (fixedKeyMap[keyName]) {
      return fixedKeyMap[keyName];
    }

    if (/id$/i.test(keyName)) {
      return () => faker.string.alphanumeric(20);
    }
    if (/phoneNumber$/i.test(keyName)) {
      return fakerJA.phone.number;
    }
    if (/^addressLine/.test(keyName)) {
      return fakerJA.location.streetAddress;
    }
    if (/url$/i.test(keyName)) {
      return faker.internet.url;
    }
  };
  `;

  const modelContents = models.flatMap(model => {
    const modelName = model.name;
    const camelModelName = camelCase(modelName);

    // nullableなJSONフィールドでDbNullへマップしないといけない
    const fieldMapLines = model.fields.reduce<string[]>((ret, f) => {
      if (!f.isRequired && f.type === 'Json') {
        const name = f.name;
        ret.push(`${name}: data.${name} === null ? Prisma.DbNull : data.${name}`);
      }
      return ret;
    }, []);

    const keyToFakePart = model.fields.reduce<string[]>((ret, f) => {
      if (f.isId) {
        // uuid
        if (f.default && typeof f.default === 'object' && 'name' in f.default && f.default.name === 'uuid(4)') {
          ret.push(`if (keyName === '${f.name}') return faker.string.uuid`);
        } else {
          ret.push(`if (keyName === '${f.name}') return () => faker.string.alphanumeric(20)`);
        }
      }

      return ret;
    }, []);

    const createParamsData = fieldMapLines.length === 0 ? `data` : `{...data, ${fieldMapLines.join(',\n')}}`;

    const func = `
export function define${modelName}Factory(db: DBClient, opts?: ZodFakeOptions) {
  const keyToFake: KeyToFake = (keyName, faker) => {
    if (opts?.keyToFake) {
      const ret = opts.keyToFake(keyName, faker);
      if (ret) return ret;
    }

    ${keyToFakePart.join('\n')}

    return commonKeyToFake(keyName, faker);
  };
  const paramsSchema = ${camelModelName}Schema.partial();
  const build = (params?: z.infer<typeof paramsSchema>, buildOpts?: ZodFakeOptions) => {
    const data = fake(${camelModelName}Schema, {faker, constraints: {int: {max: 2147483647}}, ...opts, ...buildOpts, keyToFake});
    return {...data, ...params};
  };
  const create = async (params?: z.infer<typeof paramsSchema>, buildOpts?: ZodFakeOptions) => {
    const data = build(params, buildOpts);
    return db.${camelModelName}.create({data: ${createParamsData}});
  };
  const createMany = async (paramsList: z.infer<typeof paramsSchema>[], buildOpts?: ZodFakeOptions) => {
    const dataList = paramsList.map(params => build(params, buildOpts)).map(data => (${createParamsData}));
    return db.${camelModelName}.createManyAndReturn({data: dataList});
  };
  return {
    build,
    create,
    createMany,
  };
}`;

    return func;
  });

  return [...importContents, commonKeyToFakeContent, ...modelContents].join('\n');
};

Firestoreについて

FirestoreのデータはOpenAPI specにより定義していました。

また、元々ZodスキーマもOpenAPI specから生成していました。

そのため、zod-fakeの担う部分についてはそのまま使うことができ、副次的にさほど手間をかけずにPrisma用と同様にジェネレータを実装することができました。

効果

定義されているモデルについては必ずファクトリがあり、インターフェースも統一されたことで開発生産性が高まりました。

また、GraphQLなど「全フィールドがちゃんと取得できることを確認」する際に、fillOptional オプションを使うことで非常にスッキリとテストを行うことができるようになりました。

また、良し悪しありますが、オーバーライドしていないフィールドが適当に埋まることでflakyなテストが生まれ、テストで依存しているフィールドが検知できる効果もありました。

課題

Prismaではrelationのあるモデルをネストした構造で扱うことができますが、それには対応できていません。Prismaのインターフェースが柔軟すぎるため、この対応は難しいと考えています。

一応、自動生成でなければ、ネスト構造に対応したZodスキーマを定義して zod-fake 部分を使ったファクトリーを手書きすれば少ないコストで実装することはできるので、頻出するパターンではやってもよいかもしれません。

また、特定のフィールドを固定したファクトリーを作ることは現状サポートしていないです。オーバーライドするパラメータに混ぜ込めばできるため困ってはいないですが、 defineXXXFactory 時に指定できる方がスッキリしそうです。

おわりに

明日の14日目の株式会社ビットキー Advent Calendar 2024は、Home Product チームの umi18sy が担当します。

Bitkey Developers

Discussion