Databricksにおける単体及び結合テストとテストのCI戦略

Databricksにおけるテストを考える

Databricksでテストをする記事はいくつかあるので、今回はそれをまとめます。
中でもUnit testとそこから一歩踏み込んでコード同士のdependencyを配慮するためJobの実行テストも行ってみたいと思います。
CIについても触れますが、逆に databricks bundle、デプロイ方面については触れません。

今回のゴール

Databricksでテストする方法を網羅的に探索し、それぞれの使い所、メリットデメリットをまとめてみようと思います。
また個人的な嗜好が出て恐縮ですが、pytestで単体テストは書く派なので、unittestの方は触れずにおきます。

1. 手動の実行

1.1 最も基本：Databricks Notebookによるpytest

このテストは最も基本的かつ最も簡単に実行できる内容です。

メリット

• TestするリソースもDatabricksクラスターなので本番実行されるであろうjob cluster等の環境と揃えやすい
• 面倒な初期設定が不要、いつも使っているクラスター及び環境をそのまま使える

デメリット

• 特に思いつかず

1.2 VS Code上でのpytest

今回は上記の記事の二つ目「Run tests using Databricks Connect」を使ってVS Code上で実験してみました。Databricks connectについては下記を参照してください。

これの手法を使うとVS Code上で debug&runを利用してpytestを走らせることができます。
慣れている人にとってはこちらが良いですね。

メリット

• VS code 上で debug可能、普段 Databricks connectや Extensionを利用している人はこちらでこのまま作業しやすい

デメリット

• この場合稼働しているPythonがローカルのリソースを利用するので (spark等のみリモート)、環境を完全に一致させるのは少し苦労する。

1.3 Job実行 with widgetのMode=testの活用

これは dbutils.widgets.dropdown("Mode", "Test", ["Test", "Normal"]) を利用して実行をテストモードとそうでない場合に切り替えることができます。
上記のドキュメントの中でも特にNotebookの中に関数がそのまま置かれているケースにおいて使えると記載があります。（もちろんデメリットは関数が使いまわせないことです。）
こういう場合には上記のようにDropdownで制御するだけで簡単にテストモードにできます。

また、個人的にはこれを機械学習とJob実行において利用することが多いです。

上記の画像のようにtestモードの時はデータ量とiterationを減らしてさっさと実行できます。

さらにjobの標準の変数に組み込んでおけるので、Dev, Staging環境などで走らせるjobのみこのように設定しておけばテストが簡単になります。

メリット

• 引数でテストの内容を制御できる
• Databricks jobの実行による結合テストに近い内容の制御が可能

デメリット

• コード依存、if is_testで分岐することになるので、個人の力量と思想により属人化しがちになりそう

それでは、ここまでは手動で実行し、特にCIと分離したお話でした。ここからはCIの話を進めていきたいと思います。

2. CIによるテスト実行

ここからCIの話を始めますが、正直癖がありました。今だに自分の中で最高のプラクティスは腹落ちできていません。
しかし今日はその色々と検証した結果を共有したいと思います・

2.1 (欠点あり) github actions run notebook

まず一つ目に紹介したいのが、このgithub actions run notebookです。
上記はgithub actionsのリンクと、もう一個の記事は　"Software engineering best practices for notebooks" という記事のセクション6です。
これは一応Databricksの公式ドキュメント上もBest Practiceになっているくらいです。

実際に実行してみたのが以下のスクショです。

このgithub actions何が一番素晴らしいって、実行結果の取得を待機できるんです。
みなさんもテストをCICDで組んだ人ならご経験あると思いますが、テストの結果次第でPull requestのマージを許可するか拒絶するかは自動化の面でもテストの厳密性を担保する上でも重要な検証であり、このテストの結果を得ることは上記を実現する上での最低条件であります。
しかし処理の長い結合テストや他環境に依存する必要性が高い場合はAPIでテスト用のパイプライン等をcallして、それで終了というケースも多いのではないでしょうか？
このDatabricksの提供してくれるrun-databricks-notebookはなんとテストコードから一時的なjobを作成してくれて、その結果を取得できるまで待機してくれるところも実装されている優れものなのです。
しかし。。。

github actions run notebookの残念ポイント

このtest実行、なんと結果を確認できないのです。。。
Jobを実行し終わると環境をcleanにするため、jobを削除します。
そうなると管理者レベルでしかアクセスできません。
私が試している時に最初は嘘かと思ったのですが、"Non admin users cannot view runs of local notebooks after the tmp notebook is deleted"というタイトルでissueにも上がっていました。
2023年の12月に上がっていますが、残念ながら現在対処はされていません。
対処されていない理由は色々考えられますが、多分個人的にはDatabricks Asset Bundle側で実現するの筋だと思われるので一旦こちらの開発はストップしたのではないでしょうか。多分。。。

私はあまりにもこれを使いたかったので人生初のOSSへのContibuteをしようと思いましたが、中がTypeScriptだったので、Pythonしか書けない勢はそっと閉じました。

そしてこれが理由で下記のチャレンジが続いていきます。

2.2 構築済みのjob call & 終了キャッチ

そうなってくると自前で実装する選択肢が上がってきます。
この時利用するのはDatabricks APIです。

必要なAPI endpointの htmlをClaudeに渡してサクッと書いてもらいました。

まずは自身でtestを実行するためのJobを作成しておく必要があります。
そしてそのJob idをもとにAPIでJobを動作させます。
下記は 1. Jobの実行 2. Statusを監視(ポーリング） 3. 結果が正常であれば正常終了、それ以外は異常終了するコードです。
終了はPythonのエラーが意外にも手動によるTERMINATE,CANCELやTIMEOUTもあるので注意してください。これをちゃんと書かないとgithub actionsが永遠に終わらないことになります。

import requests
import json
import os
import time
import sys

def trigger_job_run(workspace_url, token, job_id, job_parameters=None):
    """Trigger a new job run and return the run ID"""
    endpoint = f'{workspace_url}/api/2.1/jobs/run-now'
    headers = {
        'Authorization': f'Bearer {token}',
        'Content-Type': 'application/json'
    }
    
    payload = {
        'job_id': job_id
    }
    
    # Add job parameters if provided
    if job_parameters:
        payload['job_parameters'] = job_parameters
    
    print(f"Triggering job {job_id}...")
    response = requests.post(endpoint, headers=headers, json=payload)
    
    if response.status_code == 200:
        run_data = response.json()
        run_id = run_data['run_id']
        print(f"Job triggered successfully. Run ID: {run_id}")
        return run_id
    else:
        print(f'Error triggering job: {response.status_code} - {response.text}')
        sys.exit(1)

def get_job_run_status(workspace_url, token, run_id):
    """Get the current status of a job run"""
    endpoint = f'{workspace_url}/api/2.1/jobs/runs/get'
    headers = {
        'Authorization': f'Bearer {token}'
    }
    
    params = {
        'run_id': run_id
    }
    
    response = requests.get(endpoint, headers=headers, params=params)
    
    if response.status_code == 200:
        return response.json()
    else:
        print(f'Error getting run status: {response.status_code} - {response.text}')
        sys.exit(1)

def wait_for_job_completion(workspace_url, token, run_id, timeout_seconds=3600):
    """Wait for job to complete and return final status"""
    start_time = time.time()
    print(f"Waiting for job run {run_id} to complete...")
    
    while time.time() - start_time < timeout_seconds:
        run_data = get_job_run_status(workspace_url, token, run_id)
        
        state = run_data.get('state', {})
        life_cycle_state = state.get('life_cycle_state')
        result_state = state.get('result_state')
        
        print(f"Current state: {life_cycle_state}")
        
        # Check if job is complete
        if life_cycle_state in ['TERMINATED', 'SKIPPED', 'INTERNAL_ERROR']:
            return run_data
        
        # Wait before checking again
        time.sleep(30)  # Check every 30 seconds
    
    # Timeout reached
    print(f"Timeout reached after {timeout_seconds} seconds")
    return get_job_run_status(workspace_url, token, run_id)

def evaluate_job_result(run_data):
    """Evaluate job result and determine exit code"""
    state = run_data.get('state', {})
    life_cycle_state = state.get('life_cycle_state')
    result_state = state.get('result_state')
    state_message = state.get('state_message', '')
    
    print(f"\n=== Job Execution Results ===")
    print(f"Life Cycle State: {life_cycle_state}")
    print(f"Result State: {result_state}")
    print(f"State Message: {state_message}")
    
    # Print run details
    run_id = run_data.get('run_id')
    job_id = run_data.get('job_id')
    run_name = run_data.get('run_name', 'Unnamed Run')
    
    print(f"Run ID: {run_id}")
    print(f"Job ID: {job_id}")
    print(f"Run Name: {run_name}")
    
    # Print timing information if available
    start_time = run_data.get('start_time')
    end_time = run_data.get('end_time')
    if start_time and end_time:
        duration_ms = end_time - start_time
        duration_min = duration_ms / 1000 / 60
        print(f"Duration: {duration_min:.2f} minutes")
    
    # Print task results if available
    tasks = run_data.get('tasks', [])
    if tasks:
        print(f"\n=== Task Results ===")
        for task in tasks:
            task_key = task.get('task_key', 'Unknown Task')
            task_state = task.get('state', {})
            task_result = task_state.get('result_state')
            task_message = task_state.get('state_message', '')
            
            print(f"Task: {task_key}")
            print(f"  Result: {task_result}")
            if task_message:
                print(f"  Message: {task_message}")
    
    # Determine exit code based on result
    if result_state == 'SUCCESS':
        print(f"\n✅ Job completed successfully!")
        return 0
    elif result_state == 'FAILED':
        print(f"\n❌ Job failed!")
        return 1
    elif result_state == 'CANCELED':
        print(f"\n⚠️  Job was canceled!")
        return 2
    elif result_state == 'TIMEDOUT':
        print(f"\n⏰ Job timed out!")
        return 3
    elif life_cycle_state == 'SKIPPED':
        print(f"\n⏭️  Job was skipped!")
        return 4
    else:
        print(f"\n❓ Job ended with unknown state: {life_cycle_state}/{result_state}")
        return 5

def main():
    """Main execution function"""
    # Configuration
    job_id =   # Replace with your actual job ID
    token = ""
    workspace_url = ""
    
    # Optional: Add job parameters if needed
    job_parameters = {
        # "param1": "value1",
        # "param2": "value2"
    }
    
    try:        
        # Trigger job run
        run_id = trigger_job_run(workspace_url, token, job_id, job_parameters)
        
        # Wait for completion
        final_run_data = wait_for_job_completion(workspace_url, token, run_id)
        
        # Evaluate results and exit with appropriate code
        exit_code = evaluate_job_result(final_run_data)
        
        print(f"\nExiting with code: {exit_code}")
        sys.exit(exit_code)
        
    except KeyboardInterrupt:
        print("\n\n⚠️  Script interrupted by user")
        sys.exit(130)
    except Exception as e:
        print(f"\n❌ Unexpected error: {str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    main()

以下のように実行結果に依存して以上終了できます。

実行されたjobはConsoleからも同様にちゃんと確認できます。

そしてgithub actionsでこのPythonをcallすれば大丈夫です。
こんな感じでgithub actionsからも実行できます。

メリット

• 完全にDatabricks上のテストなので本番環境と同等の環境設定が容易
• jobに履歴が残るのでNotebook形式でも確認可能

デメリット

• jobに対するポーリングのコードをスクラッチで書いてしまっていて保守が必要（公式で何かあればいいのですが）
• job自体がコード上になく手動で作っているので、テストが環境依存する（※DABで解決可能・後述）

2.3 できるのか！？github actions with Databricks connect

Databricks connectを使うと、Pythonはローカル、sparkはDatabricksのSparkを利用することができます。これはSparkが使えるようになるだけでなく、Databaseへのアクセスも可能となります。

実際にこれができるかどきどきだったのですが、できましたー！
できた時のコードを共有しておきます。

こちらがgithub actions内で動作したDatabricks connectのコードです。
Python3.12、databricks-connect==15.1.0でしか動作確認できていないのでご注意ください。

if cluster_id and cluster_id.lower() != 'serverless':
    print(f"クラスター接続を使用: {cluster_id}")
    spark = (DatabricksSession.builder
            .host(f"https://{workspace_instance_name}")
            .token(token)
            .clusterId(cluster_id)
            .getOrCreate())
else:
    print(f"サーバレス接続を使用")
    spark = (DatabricksSession.builder
            .host(f"https://{workspace_instance_name}")
            .token(token)
            .serverless()
            .getOrCreate())

これによりgithub actions上でもDatabricks上のsparkを利用してtestができました。

メリット

• Spark以外Databricksの環境に依存せず簡潔
• 外部の終了を待つような処理を書かなくて良い

デメリット

• Sparkとの接続を明示する必要がある。DatabricksSession.builderなどは本来不要

2.4 pytest skipif

個人的には上記の2.2または2.3がベストプラクティスだと思っています。
以下2.4と2.5は妥協案、簡易なテスト作です。

skipifでSparkを持ちいたテストを避ける方針です。
下記のようにsparkが取得できない場合はテストを走らせない方針にすると、手動でのpytestとsparkのない環境下のテストを同時に行うコードを作ることができます。

@pytest.mark.skipif(not is_databricks_connect_available(), reason="Databricks Connect not available")
def test_get_spark():
    """Sparkセッションの取得をテスト"""
    spark = get_spark()
    assert spark is not None
    assert isinstance(spark, pyspark.sql.connect.session.SparkSession)

メリット

• 実装が簡単

デメリット

• sparkのテストを諦めている

2.5 構築済みのjob & 投げっぱなし

もう一つは 2.2の派生系でjobの投げっぱなしです。
つまりポーリングで結果の取得を諦めてjobの終了を待たないパターンです。
github actionsは自然と正常終了します。

実行されたJobはSlack/Teamsへの通知で代替し、結果の証明はスクショ等をpull requestに貼ります。
Jobをcallする程度のコードであればyamlだけでも簡単に表現できそうで管理コストもぐっと下がると思います。
また、結合テストのような長い処理の場合はgithub側で終了を待ち続けること自体がナンセンスなケースもあると思うので、状況によって利用できる場面は多いはずです。

メリット

• 実装が簡単、保守も少なめ
• 処理が長い場合にはむしろ適切

デメリット

• 終了を正しくキャッチできない可能性がある
• エラーの確認に若干運用が必要

(2.6) 多分究極は asset bundleとの組み合わせ

最後は少し感想を述べるだけですが、これらのJobの作成・テスト・削除をAsset bundleで実現できるとよりスマートになると思います。というかそれがプラクティスであることは間違い無いです。
上記の手段はjobの作成は手動で行うことが前提であり、コードとは分離した世界で環境構築する必要があるので、テスト自体が環境依存になってしまいます。

ただ、実は本質として課題に残るのはUnit testを実行した後のjobをどうするかであり、それを cleanしてなくしてしまうと、管理者しか見れないという問題は以前継続してしまいます。

うーん。。。。まあ。。。多分DABでJobを作って、それを消さずにデプロイされた状態に維持しっぱなしというのが正解かな・・・と今は思っています。

次回の記事ではそこまで発展させて踏み込んでみたいと思います。

まとめ: Databricks テスト手法まとめ（メリット / デメリット）

Databricks でのテストには多くの選択肢があり、用途やプロジェクトの規模によって最適解は変わります。
本記事で紹介したように、Notebook テストから CI 連携、Job 実行、Connect を使った外部テストまで幅広いアプローチが存在します。

それぞれの特徴を理解して使い分けることで、開発効率も品質も大きく向上できます。
あなたの環境に合った “ちょうどよいテスト戦略” を見つけるための参考になれば幸いです

手法	メリット	デメリット
1.1 Databricks Notebook で pytest	- Databricks クラスター上で実行でき本番環境と揃えやすい - 初期設定の手間がほぼ不要	- 特に大きな欠点なし
1.2 VS Code + Databricks Connect で pytest	- VS Code でデバッグしながら実行できる - 既に VS Code + Extension を使っている人は効率的	- Python はローカル実行のため、環境を完全一致させるのが難しい
1.3 Job 実行 + dbutils.widgets による Mode=test	- 引数でテスト条件を制御できる - Job 実行ベースの結合テストが簡単に行える	- if is_test などコード分岐が発生し属人化しやすい
2.1 GitHub Actions run-databricks-notebook	- Job を一時的に作成し結果が返るまで待機できる - CI/CD とテスト結果連動が可能	- 実行結果の run をユーザーが確認できない（tmp job が削除される） - 現時点で issue が未解決
2.2 独自実装：構築済み Job の call & 終了キャッチ	- Databricks 上で完全にテストを実行できる - Job に履歴が残り Notebook 形式で確認できる	- Job の状態ポーリング処理を自前で保守する必要がある
2.3 GitHub Actions + Databricks Connect（Spark Connect）	- Databricks の Spark を GitHub Actions から直接利用できる - 外部 Job を待つ処理が不要	- 接続設定（DatabricksSession.builder）が複雑になりやすい
2.4 pytest skipif（Spark がない環境ではスキップ）	- 実装が非常に簡単 - ローカル / CI 両方の pytest を一つのコードで実行できる	- Spark のテストを実質諦めている形になる
2.5 構築済み Job の投げっぱなし実行	- 実装が簡単でメンテコストが低い - 長時間の結合テストではむしろ好都合	- 結果のキャッチはできない - 正常性確認は通知や手動の運用が必要

reference

ノートブックの単体テスト
Visual Studio Code の Databricks 拡張機能を使用して Python テストを実行する
Best practices and recommended CI/CD workflows on Databricks
GitHub Actions
Databricks Connect for Python のテスト
Databricks における CI/CD
Software engineering best practices for notebooks
Databricks Connect とは
atabricks/run-notebook v0
Databricks Connect の高度な使用法