使用适用于 Visual Studio Code 的 Databricks 扩展通过 pytest 运行测试。

项目
11/11/2024

本文介绍如何使用 pytest 通过适用于 Visual Studio Code 的 Databricks 扩展运行测试。请参阅什么是 Visual Studio Code 的 Databricks 扩展？。

可以针对本地代码运行 pytest，而无需连接到远程 Azure Databricks 工作区中的群集。例如，可以使用 pytest 来测试接受 PySpark 数据帧并将其返回到本地内存的函数。若要开始使用并在本地运行 pytest，请参阅 pytest 文档中的入门。

若要针对远程 Azure Databricks 工作区中的代码运行 pytest，请在 Visual Studio Code 项目中执行以下操作：

步骤 1：创建测试

使用以下代码添加一个 Python 文件，其中包含要运行的测试。此示例假设此文件名为 spark_test.py 并且位于 Visual Studio Code 项目的根目录中。此文件包含一个 pytest 固定例程，该例程使群集的 SparkSession（群集上 Spark 功能的入口点）可用于测试。此文件包含单个测试，该测试检查表中的指定单元格是否包含指定值。你可以根据需要将自己的测试添加到此文件中。

from pyspark.sql import SparkSession
import pytest

@pytest.fixture
def spark() -> SparkSession:
  # Create a SparkSession (the entry point to Spark functionality) on
  # the cluster in the remote Databricks workspace. Unit tests do not
  # have access to this SparkSession by default.
  return SparkSession.builder.getOrCreate()

# Now add your unit tests.

# For example, here is a unit test that must be run on the
# cluster in the remote Databricks workspace.
# This example determines whether the specified cell in the
# specified table contains the specified value. For example,
# the third column in the first row should contain the word "Ideal":
#
# +----+-------+-------+-------+---------+-------+-------+-------+------+-------+------+
# |_c0 | carat | cut   | color | clarity | depth | table | price | x    | y     | z    |
# +----+-------+-------+-------+---------+-------+-------+-------+------+-------+------+
# | 1  | 0.23  | Ideal | E     | SI2     | 61.5  | 55    | 326   | 3.95 | 3. 98 | 2.43 |
# +----+-------+-------+-------+---------+-------+-------+-------+------+-------+------+
# ...
#
def test_spark(spark):
  spark.sql('USE default')
  data = spark.sql('SELECT * FROM diamonds')
  assert data.collect()[0][2] == 'Ideal'

步骤 2：创建 pytest 运行器

添加一个包含以下代码的 Python 文件，用于指示 pytest 运行在上一步骤中创建的测试。此示例假设该文件名为 pytest_databricks.py 并且位于 Visual Studio Code 项目的根目录中。

import pytest
import os
import sys

# Run all tests in the connected directory in the remote Databricks workspace.
# By default, pytest searches through all files with filenames ending with
# "_test.py" for tests. Within each of these files, pytest runs each function
# with a function name beginning with "test_".

# Get the path to the directory for this file in the workspace.
dir_root = os.path.dirname(os.path.realpath(__file__))
# Switch to the root directory.
os.chdir(dir_root)

# Skip writing .pyc files to the bytecode cache on the cluster.
sys.dont_write_bytecode = True

# Now run pytest from the root directory, using the
# arguments that are supplied by your custom run configuration in
# your Visual Studio Code project. In this case, the custom run
# configuration JSON must contain these unique "program" and
# "args" objects:
#
# ...
# {
#   ...
#   "program": "${workspaceFolder}/path/to/this/file/in/workspace",
#   "args": ["/path/to/_test.py-files"]
# }
# ...
#
retcode = pytest.main(sys.argv[1:])

步骤 3：创建自定义运行配置

若要指示 pytest 运行测试，必须创建自定义运行配置。使用现有的基于 Databricks 群集的运行配置来创建你自己的自定义运行配置，如下所示：

在主菜单中，单击“运行”>“添加配置”。
在“命令面板”中，选择“Databricks”。

Visual Studio Code 会向项目添加一个 .vscode/launch.json 文件（如果该文件尚不存在）。
如下所示更改启动器运行配置，然后保存文件：
- 将此运行配置的名称从 Run on Databricks 更改为此配置的唯一显示名称（在本示例中为 Unit Tests (on Databricks)）。
- 将 program 从 ${file} 更改为项目中包含测试运行器的路径（在本示例中为 ${workspaceFolder}/pytest_databricks.py）。
- 将 args 从 [] 更改为项目中包含测试文件的路径（在本例中为 ["."]）。
launch.json 文件应如下所示：
```
{
  // Use IntelliSense to learn about possible attributes.
  // Hover to view descriptions of existing attributes.
  // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
  "version": "0.2.0",
  "configurations": [
    {
      "type": "databricks",
      "request": "launch",
      "name": "Unit Tests (on Databricks)",
      "program": "${workspaceFolder}/pytest_databricks.py",
      "args": ["."],
      "env": {}
    }
  ]
}
```

步骤 4：运行测试

首先确保 pytest 已安装在群集上。例如，在 Azure Databricks 工作区中打开群集的设置页后，执行以下操作：

在“库”选项卡上，如果“pytest”可见，则表示已安装 pytest。如果“pytest”不可见，请单击“安装新程序”。
对于“库源”，请单击“PyPI”。
对于“包”，请输入 pytest。
单击“安装” 。
等待“状态”从“挂起”更改为“已安装”。

若要运行测试，请在 Visual Studio Code 项目中执行以下操作：

在主菜单中，单击“查看”>“运行”。
在“运行和调试”列表中，单击“单元测试(在 Databricks 上)”（如果尚未选择）。
单击绿色箭头（“开始调试”）图标。

pytest 结果将显示在“调试控制台”（在主菜单中选择“查看”>“调试控制台”）中。例如，这些结果显示至少在 spark_test.py 文件中找到了一个测试，句点 (.) 表示找到并通过了一项测试。（失败的测试将显示为 F。）

<date>, <time> - Creating execution context on cluster <cluster-id> ...
<date>, <time> - Synchronizing code to /Workspace/path/to/directory ...
<date>, <time> - Running /pytest_databricks.py ...
============================= test session starts ==============================
platform linux -- Python <version>, pytest-<version>, pluggy-<version>
rootdir: /Workspace/path/to/directory
collected 1 item

spark_test.py .                                                          [100%]

============================== 1 passed in 3.25s ===============================
<date>, <time> - Done (took 10818ms)

通过

使用适用于 Visual Studio Code 的 Databricks 扩展通过 pytest 运行测试。

步骤 1：创建测试

步骤 2：创建 pytest 运行器

步骤 3：创建自定义运行配置

步骤 4：运行测试

其他资源