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

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

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

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

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

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

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

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 ランナーを作成する

前のステップからテストを実行するようにpytestに指示する次のコードを含むPythonファイルを追加します。 この例では、ファイルの名前が 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 は既にインストールされています。 pytest が表示されない場合は、[新規インストール] をクリックします。
  2. [ライブラリ ソース] で、[PyPI] をクリックします。
  3. [パッケージ] に「pytest」と入力します。
  4. [インストール] をクリックします。
  5. [状態] が [保留中] から [インストール済み] に変わるまで待ちます。

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

  1. メイン メニューで、[ビュー] > [実行] をクリックします。
  2. まだ選択していない場合は、[実行とデバッグ] の一覧で、[Unit Tests (on Databricks)](単体テスト (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 API を使用するテストをローカルで実行するには、Databricks Connect を使用します。

手順 1: Databricks Connect を構成する

手順に従って、拡張機能の Databricks Connect を構成します。 Visual Studio Code の Databricks 拡張機能の Databricks Connect を使用した Debug コードを参照してください。

手順 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 起動構成を追加または更新する

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

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

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

    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