SQLMesh 测试自动化：提升数据工程效率

在现代数据工程中，确保数据模型的准确性和可靠性至关重要。SQLMesh 提供了一套强大的测试工具，用于验证数据模型的输出是否符合预期。本文将深入探讨 SQLMesh 的测试功能，包括如何创建测试、支持的数据格式以及如何运行和调试测试。

SQLMesh 测试概述

SQLMesh 的测试功能旨在通过持续验证每个模型的输出来保护项目免受回归影响。与软件开发的单元测试类似，SQLMesh 使用预定义的输入评估模型的逻辑，并将其输出与每个测试提供的预期结果进行比较。这种测试方法不仅可以在每次新计划创建时自动执行，还可以作为 CI/CD 流程的一部分按需执行。

创建测试

在 SQLMesh 中，测试套件是一个包含在项目 tests/ 文件夹中的 YAML 文件，文件名以 test 开头并以 .yaml 或 .yml 结尾。测试套件可以包含一个或多个唯一命名的单元测试，每个单元测试都有一系列属性来定义其行为。一个单元测试至少需要指定被测试的模型、上游模型的输入值以及目标模型的查询和/或公共表表达式的预期输出。

SQLMesh 支持多种方式来定义单元测试中的输入和输出数据：

YAML 字典：列映射到它们的值。
CSV：逗号分隔的值。
SQL 查询：针对测试连接执行 SQL 查询以生成数据。

入门实例

在本例中，我们将使用sqlmesh_example.Full_model模型，作为sqlmesh init命令的一部分提供，定义如下：

sql 复制代码

MODEL (
  name sqlmesh_example.full_model,
  kind FULL,
  cron '@daily',
  grain item_id,
  audits (assert_positive_order_ids),
);

SELECT
  item_id,
  COUNT(DISTINCT id) AS num_orders,
FROM
  sqlmesh_example.incremental_model
GROUP BY item_id

此模型从上游的 sqlmesh_example.incremental_model 中聚合每个 item_id 的订单数量。测试此模型的一种方法如下所示：

yml 复制代码

test_example_full_model:
  model: sqlmesh_example.full_model
  inputs:
    sqlmesh_example.incremental_model:
      rows:
      - id: 1
        item_id: 1
      - id: 2
        item_id: 1
      - id: 3
        item_id: 2
  outputs:
    query:
      rows:
      - item_id: 1
        num_orders: 2
      - item_id: 2
        num_orders: 1

此测试验证 sqlmesh_example.full_model 是否能正确统计每个 item_id 的订单数量。它向 sqlmesh_example.incremental_model 提供三行输入，并期望目标模型的查询输出两行结果。

运行和调试测试

SQLMesh 的测试可以通过 CLI 或 Jupyter 笔记本按需执行。CLI 命令 sqlmesh test 可以用来执行所有测试，而 %run_test 笔记本魔法命令则允许在笔记本环境中执行测试。如果遇到测试失败，可以使用 --preserve-fixtures 选项保留输入夹具，以便进行调试。

shell 复制代码

$ sqlmesh test
.
----------------------------------------------------------------------
Ran 1 test in 0.005s

OK

要运行特定的模型测试，请传入测试套件文件名后跟：：和测试名称：

sh 复制代码

$ sqlmesh test tests/test_full_model.yaml::test_example_full_model

您还可以使用通配符路径扩展语法运行匹配模式或子字符串的测试：

sh 复制代码

$ sqlmesh test tests/test_*

测试用例实战

CTE测试

模型查询中的各个公用表表达式（CTE）也可以进行测试。为了演示这一点，让我们对 sqlmesh_example.full_model 的查询稍作修改，添加一个名为 filtered_orders_cte 的 CTE：

sql 复制代码

WITH filtered_orders_cte AS (
  SELECT
    id,
    item_id
  FROM
    sqlmesh_example.incremental_model
  WHERE
    item_id = 1
)
SELECT
  item_id,
  COUNT(DISTINCT id) AS num_orders,
FROM
  filtered_orders_cte
GROUP BY item_id

下面的测试将在聚合发生之前验证该CTE的输出：

yml 复制代码

test_example_full_model:
  model: sqlmesh_example.full_model
  inputs:
    sqlmesh_example.incremental_model:
        rows:
        - id: 1
          item_id: 1
        - id: 2
          item_id: 1
        - id: 3
          item_id: 2
  outputs:
    ctes:
      filtered_orders_cte:
        rows:
          - id: 1
            item_id: 1
          - id: 2
            item_id: 1
    query:
      rows:
      - item_id: 1
        num_orders: 2

csv文件

这就是我们如何定义与第一个示例相同的测试，但输入数据格式为CSV：

yml 复制代码

test_example_full_model:
  model: sqlmesh_example.full_model
  inputs:
    sqlmesh_example.incremental_model:
      format: csv
      rows: |
        id,item_id
        1,1
        2,1
        3,2
  outputs:
    query:
      rows:
      - item_id: 1
        num_orders: 2
      - item_id: 2
        num_orders: 1

sql查询

这就是我们如何能够将上述第一个示例中的相同测试定义为这样一种形式，只不过输入数据是通过 SQL 查询生成的：

yml 复制代码

test_example_full_model:
  model: sqlmesh_example.full_model
  inputs:
    sqlmesh_example.incremental_model:
      query: |
        SELECT 1 AS id, 1 AS item_id
        UNION ALL
        SELECT 2 AS id, 1 AS item_id
        UNION ALL
        SELECT 3 AS id, 2 AS item_id
  outputs:
    query:
      rows:
      - item_id: 1
        num_orders: 2
      - item_id: 2
        num_orders: 1

数据文件

SQLMesh支持从外部文件加载数据。要实现这一点，你可以使用pathattribute，它指定要加载的数据的路径名：

yml 复制代码

test_example_full_model:
  model: sqlmesh_example.full_model
  inputs:
    sqlmesh_example.incremental_model:
      format: csv
      path: filepath/test_data.csv

如果省略format，则该文件将作为YAML文档加载。

省略列

对于宽表（即具有众多列的表），定义完整的输入和预期输出可能会变得繁琐。因此，如果某些列可以安全地忽略，那么它们可以从任何行中省略，并且对于该行，其值将被视为 NULL。

此外，可以通过将 partial 设置为 true 来仅测试感兴趣的输出列的一部分：

yml 复制代码

  outputs:
    query:
      partial: true
      rows:
        - <column_name>: <column_value>
          ...

当缺失的列不能被视为 NULL 值，但我们仍希望忽略这些列时，此设置非常有用。若要将此设置应用于所有预期输出，请在"输出"键下进行设置：

yml 复制代码

  outputs:
    partial: true
    ...

冻结时间

某些模型可能会使用计算给定时间点 datetime 值的 SQL 表达式，例如 CURRENT_TIMESTAMP。由于这些表达式是非确定性的，仅仅指定预期的输出值不足以对其进行测试。

通过设置 execution_time 宏变量来模拟测试上下文中的当前时间解决了这个问题，从而使其值具有确定性。

以下示例展示了如何使用 execution_time 来测试使用 CURRENT_TIMESTAMP 计算的列。我们将要测试的模型定义如下：

sql 复制代码

MODEL (
  name colors,
  kind FULL
);

SELECT
  'Yellow' AS color,
  CURRENT_TIMESTAMP AS created_at

测试文件如下：

yaml 复制代码

test_colors:
  model: colors
  outputs:
    query:
      - color: "Yellow"
        created_at: "2023-01-01 12:05:03"
  vars:
    execution_time: "2023-01-01 12:05:03"

还可以为执行时间设置时区，方法是在时间戳字符串中包含该时区。

如果提供了时区，目前的要求是测试的预期日期时间值必须是无时区的时戳，这意味着它们需要相应地进行偏移。

如果我们希望将时间冻结为 UTC+2，以下是上述测试的编写方式：

yaml 复制代码

test_colors:
  model: colors
  outputs:
    query:
      - color: "Yellow"
        created_at: "2023-01-01 10:05:03"
  vars:
    execution_time: "2023-01-01 12:05:03+02:00"

自动生成测试

手动创建测试可能会显得单调乏味且容易出错，这就是为什么 SQLMesh 还提供了使用 create_test 命令来实现自动化处理这一过程的方法。

此命令能够为给定的模型生成完整的测试，只要其上游模型的表存在于项目的数据仓库中，并且这些表中已有数据即可。

实战示例

在这个示例中，我们将展示如何为 sqlmesh_example.incremental_model 生成测试用例。sqlmesh_example.incremental_model 是作为 sqlmesh init 命令的一部分提供的另一个模型，其定义如下：

sql 复制代码

MODEL (
  name sqlmesh_example.incremental_model,
  kind INCREMENTAL_BY_TIME_RANGE (
    time_column event_date
  ),
  start '2020-01-01',
  cron '@daily',
  grain (id, event_date)
);

SELECT
  id,
  item_id,
  event_date,
FROM
  sqlmesh_example.seed_model
WHERE
  event_date BETWEEN @start_date AND @end_date

首先，我们需要明确上游模型 sqlmesh_example.seed_model 的输入数据。create_test 命令的执行始于对项目的数据仓库发出用户自定义的查询，以获取这些数据。

例如，以下查询将从与模型 sqlmesh_example.seed_model 相对应的表中返回三行数据：1

sql 复制代码

SELECT * FROM sqlmesh_example.seed_model LIMIT 3

接下来，请留意 sqlmesh_example.incremental_model 中包含一个引用了 @start_date 和 @end_date 宏变量的过滤条件。为了使生成的测试具有确定性，从而确保它总是能够成功，我们需要定义这些变量，并修改上述查询以相应地约束 event_date。

如果我们将 @start_date 设为 '2020-01-01' 并将 @end_date 设为 '2020-01-04'，上述查询需要修改为：

sql 复制代码

SELECT * FROM sqlmesh_example.seed_model WHERE event_date BETWEEN '2020-01-01' AND '2020-01-04' LIMIT 3

运行此操作会创建以下新测试，其位于 tests/test_incremental_model.yaml 文件中：

sql 复制代码

test_incremental_model:
  model: sqlmesh_example.incremental_model
  inputs:
    sqlmesh_example.seed_model:
    - id: 1
      item_id: 2
      event_date: 2020-01-01
    - id: 2
      item_id: 1
      event_date: 2020-01-01
    - id: 3
      item_id: 3
      event_date: 2020-01-03
  outputs:
    query:
    - id: 1
      item_id: 2
      event_date: 2020-01-01
    - id: 2
      item_id: 1
      event_date: 2020-01-01
    - id: 3
      item_id: 3
      event_date: 2020-01-03
  vars:
    start: '2020-01-01'
    end: '2020-01-04'

配置测试连接

对于给定的测试，可以更改测试连接。例如，当被测试的模型无法正确编译为默认测试引擎的方言时，这可能会很有用。

以下示例通过修改 test_example_full_model 来演示这一点，使其针对单线程本地 Spark 进程运行，该进程在项目的 config.yaml 文件中的 spark_testing 网关中定义为 test_connection：

复制代码

gateways:
  local:
    connection:
      type: duckdb
      database: db.db
  spark_testing:
    test_connection:
      type: spark
      config:
        # Run Spark locally with one worker thread
        "spark.master": "local"

        # Move data under /tmp so that it is only temporarily persisted
        "spark.sql.warehouse.dir": "/tmp/data_dir"
        "spark.driver.extraJavaOptions": "-Dderby.system.home=/tmp/derby_dir"

default_gateway: local

model_defaults:
  dialect: duckdb

修改测试用例：

yml 复制代码

test_example_full_model:
  gateway: spark_testing
  # ... the other test attributes remain the same

最后总结

SQLMesh 的测试功能为数据工程师提供了一个强大的工具，用于确保数据模型的准确性和可靠性。通过自动化测试过程，SQLMesh 帮助团队在每次模型变更后都能快速验证其正确性。无论是手动创建测试还是使用自动测试生成工具，SQLMesh 都能有效地提升数据工程的质量和效率。对于希望在数据工程中实现更高可靠性的团队来说，SQLMesh 是一个不可或缺的工具。