使用 Databricks IDE 扩展运行 Python 测试

本页介绍如何使用 Databricks IDE 扩展运行 Python 测试。 参见 Databricks的IDE扩展。

使用 pytest 运行测试

可以针对本地代码运行 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 群集的运行配置来创建你自己的自定义运行配置，如下所示：

  1. 在主菜单中，单击“运行”>“添加配置”。

  2. 在“命令面板”中，选择“Databricks”。

    Visual Studio Code 会向项目添加一个 .vscode/launch.json 文件（如果该文件尚不存在）。

  3. 如下所示更改启动器运行配置，然后保存文件：

    • 将此运行配置的名称从 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 工作区中打开群集的设置页后，执行以下操作：

  1. 在“库”选项卡上，如果“pytest”可见，则表示已经安装。 如果“pytest”不可见，请单击“安装新程序”。

  2. 对于“库源”，请单击“PyPI”。

  3. 对于包，请输入pytest。

  4. 单击“安装” 。

  5. 等待“状态”从“挂起”更改为“已安装”。

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

  1. 在主菜单中，单击“查看”>“运行”。

  2. 在“运行和调试”列表中，单击“单元测试(在 Databricks 上)”（如果尚未选择）。

  3. 单击绿色箭头（“开始调试”）图标。

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)

使用 Databricks Connect 运行测试

若要在本地运行使用 Spark API 的测试，请使用 Databricks Connect。

步骤 1：配置 Databricks Connect

遵循以下步骤配置用于扩展的 Databricks Connect。 请参阅 使用 Databricks Connect 为 Databricks IDE 扩展调试代码。

步骤 2：创建单元测试

使用以下代码添加 Python 文件，其中包含要运行的测试。 此示例假定此文件已命名 main_test.py。

from my_project import main

def test_find_all_taxis():
    taxis = main.find_all_taxis()
    assert taxis.count() > 5

步骤 3：添加或更新 debugpy 启动配置

接下来，创建一个 debugpy 启用 Databricks Connect 的启动配置。

  1. 在 Visual Studio Code 的主菜单上，单击“ 运行 > 添加配置”。

  2. 在 命令面板中，选择 “Python 调试器”。

    Visual Studio Code 会向项目添加一个 .vscode/launch.json 文件（如果该文件尚不存在）。

  3. 添加 "databricks": true 字段。 这将启用 Databricks Connect。

{
  "version": "0.2.0",
  "configurations": [
    {
      "name": "Unit Tests (on Databricks)",
      "type": "debugpy",
      "databricks": true,
      "request": "launch",
      "program": "${file}",
      "args": ["."],
      "env": {},
      "console": "integratedTerminal"
    }
  ]
}

步骤 4：运行测试

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

  1. 在主菜单上，单击“ 查看 > 测试 ”以打开测试面板。
  2. 在测试面板中，单击与之 main_test.py关联的调试图标来运行测试。 请注意，仅运行测试不会触发修改后的调试配置，并且代码无法访问 Databricks Connect。
  • Last updated on