Visual Studio Code の Databricks 拡張機能を使用して pytest でテストを実行する
この記事では、 pytest
を使用してテストを実行する方法と、Visual Studio Code の Databricks 拡張機能を使用する方法について説明します。 「Visual Studio Code の Databricks 拡張機能とは」を参照してください。
pytest は、リモート Databricks ワークスペース内のクラスターへの接続を必要としないローカル コードで実行できます。たとえば、 pytest
を使用して、ローカル メモリ内の PySpark DataFrames を受け入れて返す関数をテストできます。 pytest
の使用を開始し、ローカルで実行するには、pytest
のドキュメントの「はじめに」を参照してください。
リモート Databricks ワークスペースのコードで pytest
を実行するには、Visual Studio Code プロジェクトで次の操作を行います。
ステップ 1: テストを作成する
実行するテストを含む次のコードを含む Python ファイルを追加します。 この例では、このファイルの名前が spark_test.py
で、Visual Studio Code プロジェクトのルートにあることを前提としています。 このファイルには、クラスターのSparkSession
(クラスター上の Spark 機能へのエントリ ポイント) をテストで使用できるようにするための pytest
フィクスチャが含まれています。このファイルには、テーブル内の指定したセルに指定した値が含まれているかどうかを確認する 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 ランナーを作成する
次のコードを含む 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
がクラスターに既にインストールされていることを確認してください。 たとえば、Databricks ワークスペースでクラスターの設定ページを開いた状態で、次の操作を行います。
「ライブラリ」タブで、pytest が表示されている場合、
pytest
はすでにインストールされています。pytest が表示されない場合は、[新規インストール] をクリックします。[ライブラリ ソース] で [PyPI] をクリックします。
[パッケージ] に「
pytest
」と入力します。[インストール] をクリックします。
ステータスが [保留中] から [インストール済み] に変わるまで待ちます。
テストを実行するには、Visual Studio Code プロジェクトから次の操作を行います。
メイン メニューで、[ 表示] > [実行] をクリックします。
[ 実行とデバッグ] リストで、[ 単体テスト (Databricks)] をクリックします (まだ選択されていない場合)。
緑色の矢印 (デバッグの開始) アイコンをクリックします。
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)