メインコンテンツまでスキップ

Visual Studio Code の Databricks 拡張機能を使用して Python テストを実行する

このページでは、Visual Studio Code の Databricks 拡張機能を使用して Python テストを実行する方法について説明します。「Visual Studio Code の Databricks 拡張機能とは何ですか?」を参照してください。

pytestを使用してテストを実行する​

pytest は、リモート Databricks ワークスペース内のクラスターへの接続を必要としないローカル コードで実行できます。たとえば、 pytest を使用して、ローカル メモリ内の PySpark データフレーム を受け入れて返す関数をテストできます。 pytestの使用を開始し、ローカルで実行するには、pytest のドキュメントの「はじめに」を参照してください。

リモート Databricks ワークスペースのコードで pytest を実行するには、Visual Studio Code プロジェクトで次の操作を行います。

ステップ1: テストを作成する​

実行するテストを含む次のコードを含む Python ファイルを追加します。この例では、このファイルの名前が spark_test.py で、Visual Studio Code プロジェクトのルートにあることを前提としています。pytestこのファイルには、クラスターのSparkSession Spark(クラスターの 機能へのエントリ ポイント) をテストで使用できるようにするための フィクスチャ が含まれています。このファイルには、テーブル内の指定したセルに指定した値が含まれているかどうかを確認する 1 つのテストが含まれています。必要に応じて、このファイルに独自のテストを追加できます。

Python
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 プロジェクトのルートにあることを前提としています。

Python
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ファイルは次のようになります。

    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 がクラスターに既にインストールされていることを確認してください。 たとえば、 Databricks ワークスペースでクラスターの設定ページを開いた状態で、次の操作を行います。

  1. 「ライブラリ」 タブで、 pytest が表示されている場合、pytest はすでにインストールされています。 pytest が表示されない場合は、[ 新規インストール ] をクリックします。
  2. [ライブラリ ソース ] で [ PyPI] をクリックします。
  3. [パッケージ ] に「pytest」と入力します。
  4. [ インストール ] をクリックします。
  5. ステータスが「 保留中 」から「 インストール済み 」に変わるまで待ってください。

テストを実行するには、Visual Studio Code プロジェクトから次の操作を行います。

  1. メイン メニューで、[ 表示] > [実行 ] をクリックします。
  2. [ 実行とデバッグ] リストで、[ 単体テスト (Databricks)] をクリックします (まだ選択されていない場合)。
  3. 緑色の矢印 ( デバッグの開始 ) アイコンをクリックします。

pytest結果は 、デバッグコンソール (メインメニューの 「デバッグコンソールの表示」> )に表示されます。たとえば、これらの結果は、 spark_test.py ファイルで少なくとも 1 つのテストが見つかったことを示しており、ドット (.) は 1 つのテストが見つかり、合格したことを意味します。 (テストに失敗すると、 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 APIs使用するテストをローカルで実行するには、 Databricks Connect使用します。

ステップ 1: Databricks Connectを構成する​

ステップに従って、拡張機能用にDatabricks Connect構成します。 Visual Studio Code の Databricks 拡張機能については、「Databricks Connect を使用してコードをデバッグする」を参照してください。

ステップ 2: 単体テストを作成する​

実行するテストを含む次のコードを含む Python ファイルを追加します。この例では、このファイルの名前がmain_test.pyであると想定しています。

Python
from my_project import main


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

ステップ 3: debugpy 起動構成を追加または更新します​

次に、Databricks Connect を有効にするdebugpy起動構成を作成します。

  1. Visual Studio Code のメイン メニューで、 [実行] > [構成の追加] をクリックします。

  2. コマンドパレット で、 Python デバッガー を選択します。

    Visual Studio Code は、 .vscode/launch.json ファイルをプロジェクトに追加します (このファイルがまだ存在しない場合)。

  3. "databricks": trueフィールドを追加します。これにより、Databricks Connect が有効になります。

JSON
{
  "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 にアクセスできないことに注意してください。
に最終更新