トップ > BigQuery > BigQueryを使ったロジックのテストについて

BigQueryを使ったロジックのテストについて

BigQuery Node.js Testing

ogp

はじめに

こんにちは、enechain で eScan という ETRM の開発を担当している Yamato です。 ETRM とは、エネルギー取引にまつわるリスクを管理するソフトウェアで、電力の調達データ・販売量データ・電力市場の価格・燃料価格などを元にユーザが抱えるリスクを可視化し管理できるツールです。

リスクの算出には、大規模なデータをもとに高度な計算を行う必要があります。 eScan では開発当初、Node.js 上で計算を行っておりましたが、より高速で安定した処理が行えるよう、BigQuery に計算基盤を移行しました。

techblog.enechain.com

BigQuery を活用することで、目論見どおり計算の高速化・安定化が図れました。

一方で、BigQuery に移行したことでどのようにテストを行うかが課題となりました。 当然ながら計算を行うクエリは eScan のコアであり、ロジックの担保やデグレーションの防止などの観点からテストを実装することが不可欠です。

ただ、BigQuery に対するテストは標準となる手法が確立されていません。 eScan では手探りながら BigQuery へのテストの方法を構築してきましたので、その現状を紹介したいと思います。

前提

eScan は、フレームワークとしてNestJSを採用しています。また、テストフレームワークとしてVitestを用いています。

また、計算のクエリは基本的にはほぼ全て以下のような流れになっています。

  1. データをかき集める
  2. 計算する
  3. 結果をテーブルに保存する
BEGIN
    INSERT INTO
        datasetId.SummarizedRevenue (
        area,
        yearMonth,
        volume)
    SELECT
      area,
      DATE_TRUNC(date, MONTH) as yearMonth,
      SUM(volume) AS volume
    FROM datasetId.RawRevenue
    GROUP BY area, DATE_TRUNC(date, MONTH);
END

「保存された契約から販売量を取り出す」「販売量と単価をかけて売上を算出する」など、ある程度の粒度でクエリが分割し、次のクエリがその結果を使って処理をおこなっていく、といったイメージです。

また、クエリ自体はインフラ層に閉じ込めており、各計算クエリを以下のように実行しています。

export class CalcRevenue {
  constructor(private readonly bqClient: BqClient) {}

  async executeQuery(query: Query): Promise<QueryRowsResponse> {
    const [job, _] = await bqClient.dataset().createQueryJob(query);

    return job.getQueryResults();
  }

  async calc() {
    const queryString = readFileSync(
      resolve(__dirname, "計算対象のクエリ.sql")
    ).toString();

    const query = {
      query: queryString,
    };

    await this.executeQuery(organizationId, query);
  }
}

方針

基本的な考え

クエリの実行自体をインフラ層に閉じ込めているため、テストはインフラ層のメソッドに対するテストと言う形を取ります。 以下のような流れでテストを落とし込んでいきます。

  1. 必要なデータを insert
  2. クエリを実行
  3. 実行結果を結果テーブルから select して検証

テストは vitest で実行するため、特別なことをしなくてもローカル環境や CI 等、既存の環境で実行できます。

クエリ自体をどこで実行するか、という点については現在のところ 2 つの方法を取っています。

1 つ目が BigQuery Emulator を使ってテストを実行する方法です。 eScan では、大半のクエリがこの Emulator を用いてテストを実行しています。 利点は、なんといってもコストやセキュリティの面を気にせずにテストを実行できることです。 公式のエミュレータが提供されていない中、OSS でありながらほとんどの場面で問題なく使えるため、テストを行う際には非常に有用です。

BigQuery Emulator に関する詳細は以下の記事をご覧ください。 techblog.enechain.com

2 つ目が、JavaScript UDF を使っている箇所など Emulator に実装されていない機能を利用したクエリは、実際に BigQuery を叩けるようにしておくという方法です。 上記の通り Emulator は強力なツールですが、所々で BigQuery と異なる挙動をする場合や、サポートされていない関数などがあります。 そのようなクエリに対するテストは、実際の BigQuery をたたくような仕組みを用意してカバーしています。

実際のコード例

以上のような点をふまえたうえで、プリミティブな実装例が以下になります。

describe("calc", () => {
  let calcRevenue: CalcRevenue;
  let bqClient: BqClient;

  beforeAll(async () => {
    const app: TestingModule = await Test.createTestingModule({
      providers: [BqClient, CalcRevenue],
      imports: [CustomConfigModule],
    }).compile();
    calcRevenue = app.get<CalcRevenue>(CalcRevenue);
    bqClient = app.get<BqClient>(BqClient);

    // テスト用のdatasetを作成する
    await bqClient.createDataset(datasetId, {
      location: "asia-northeast1",
    });

    // 前提となるテーブルと結果が格納されるテーブルを作成しておく
    await bqClient.migrate(rawRevenueSchema, summarizedRevenueSchema);

    // テスト用のデータを作成する
    await bqClient.dataset().createQueryJob(`
        INSERT INTO \` ${datasetId}.RawRevenue\` (
            area,
            date,
            volume)
        VALUES
            (
                'Tokyo',
                '2023-08-01',
                1),
            (
                'Tokyo',
                '2023-08-02',
                2),
            (
                'Tokyo',
                '2023-08-03',
                3),
            (
                'Tokyo',
                '2023-08-04',
                4),
    `);
  });

  afterAll(async () => {
    // テスト用のdatasetを削除する
    await bqClient.dataset().delete({ force: true });
  });

  test("2023/08の値が正しいこと", async () => {
    await calcRevenue.calc();
    const [job, _] = await bqClient.dataset().createQueryJob(`
        SELECT *
        FROM ${datasetId}.SummarizedRevenue
        WHERE
            yearMonth = '2023-08-01'
    `);

    const [result] = await job.getQueryResults();

    expect(result[0].volume).toEqual(10);
  });
});

テーブルの作成等については下記の記事にまとめております。

techblog.enechain.com

工夫点

Emulator と BigQuery の向き先切り替え

上述の通り、クエリの実行は Emulator と BigQuery の両方を扱っています。 そのためテスト毎に切り替えをできる仕組みを用意しておくと便利です。 Google 公式の BigQuery SDK のクラスををラップする形で Client クラスを用意し、initialize 時に接続の向き先を Emulator か BigQuery か指定できるようにしています。

export class BqClient extends BigQuery {
  constructor(
    private readonly toEmulator: boolean
  ) {
     const config = toEmulator
      ? {
          ...BQ_CONFIG
          projectId: EMULATOR_PROJECT_ID
          apiEndpoint: EMULATOR_END_POINT
        }
      : BQ_CONFIG
    super(config)
  }
  ...
}

テスト毎に Dataset を作り終了したら破棄する

テスト時に Insert されたデータはどのように後始末すべきでしょうか。 Rspec のuse_transactional_fixturesのように、テスト中のトランザクションを隔離してデータのロールバックを自動で行うことは難しいです。

そこで、テスト毎に Dataset を新規に create、終了したらAfterAllで Dataset ごと破棄し、常にクリーンな状態でテストが実行できるようにしています。 ただし、カレンダーテーブルなど書き込みが発生しないようなマスタテーブルは、別の Dataset に隔離して削除しないようにしています。

test の timeout を例外的に延長する

Emulator であれ BigQuery であれ、どちらもテストが遅いです。 Vitest の timeout はデフォルトで 5 秒ですが、実際には 1 つのケースで 10 秒前後かかる場合もあり到底足りません。 テストが遅いこと自体改善すべきポイントではあるのですが、タイムアウトを延長することで一旦の対応を行っています。

vitest の場合は、vitest.config.js でタイムアウトを切り替えられるようにしておくと良いでしょう。

export default defineConfig({
  test: {
    testTimeout: process.env.BIG_QUERY_TEST === "on" ? 60 * 1000 : 5000,
  },
});

まとめ

この記事では eScan における BigQuery のテストについてご紹介いたしました。 もちろんこれが完成ではなく、試行錯誤を行いながら、より安全で使いやすい状態を目指して引き続きを改善を行ってまいります。

enechain では、プロダクトを一緒に創っていく仲間を募集しています。少しでもご興味・ご関心がございましたら、ぜひお気軽に以下からご応募よろしくお願いします。

herp.careers