Execute testes com pytest usando a extensão Databricks para Visual Studio Code

  • Artigo

Este artigo descreve como executar testes usando pytest a extensão Databricks para Visual Studio Code. Consulte O que é a extensão Databricks para Visual Studio Code?.

Você pode executar pytest em código local que não precisa de uma conexão com um cluster em um espaço de trabalho remoto do Azure Databricks. Por exemplo, você pode usar pytest para testar suas funções que aceitam e retornam DataFrames PySpark na memória local. Para começar pytest e executá-lo localmente, consulte Introdução na pytest documentação.

Para executar pytest em código em um espaço de trabalho remoto do Azure Databricks, faça o seguinte em seu projeto do Visual Studio Code:

Etapa 1: Criar os testes

Adicione um arquivo Python com o código a seguir, que contém seus testes a serem executados. Este exemplo pressupõe que esse arquivo é nomeado spark_test.py e está na raiz do seu projeto Visual Studio Code. Esse arquivo contém um pytest acessório, que torna o cluster SparkSession (o ponto de entrada para a funcionalidade do Spark no cluster) disponível para os testes. Esse arquivo contém um único teste que verifica se a célula especificada na tabela contém o valor especificado. Você pode adicionar seus próprios testes a este arquivo conforme necessário.

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'

Etapa 2: Criar o corredor pytest

Adicione um arquivo Python com o código a seguir, que instrui pytest a executar seus testes da etapa anterior. Este exemplo pressupõe que o arquivo é nomeado pytest_databricks.py e está na raiz do seu projeto 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:])

Etapa 3: Criar uma configuração de execução personalizada

Para instruir pytest a executar seus testes, você deve criar uma configuração de execução personalizada. Use a configuração de execução baseada em cluster Databricks existente para criar sua própria configuração de execução personalizada, da seguinte maneira:

  1. No menu principal, clique em Executar > Adicionar configuração.

  2. Na Paleta de comandos, selecione Databricks.

    O Visual Studio Code adiciona um .vscode/launch.json arquivo ao seu projeto, se esse arquivo ainda não existir.

  3. Altere a configuração de execução inicial da seguinte maneira e salve o arquivo:

    • Altere o nome desta configuração de execução de Run on Databricks para algum nome de exibição exclusivo para esta configuração, neste exemplo Unit Tests (on Databricks).
    • Mude program de para o caminho no projeto que contém o executor de ${file} teste, neste exemplo ${workspaceFolder}/pytest_databricks.py.
    • Mude args de [] para o caminho no projeto que contém os arquivos com seus testes, neste exemplo ["."].

    O seu launch.json ficheiro deve ter o seguinte aspeto:

    {
  // 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": {}
    }
  ]
}

Etapa 4: Executar os testes

Certifique-se de que pytest já está instalado no cluster primeiro. Por exemplo, com a página de configurações do cluster aberta em seu espaço de trabalho do Azure Databricks, faça o seguinte:

  1. Na guia Bibliotecas, se pytest estiver visível, então pytest já está instalado. Se pytest não estiver visível, clique em Instalar novo.
  2. Para Library Source, clique em PyPI.
  3. Em Pacote, digite pytest.
  4. Clique em Install (Instalar).
  5. Aguarde até que o Status mude de Pendente para Instalado.

Para executar os testes, faça o seguinte a partir do seu projeto Visual Studio Code:

  1. No menu principal, clique em Exibir > Executar.
  2. Na lista Executar e Depurar, clique em Testes de Unidade (no Databricks), se ainda não estiver selecionado.
  3. Clique no ícone de seta verde (Iniciar Depuração).

Os pytest resultados são exibidos no Console de Depuração (Exibir > Console de Depuração no menu principal). Por exemplo, esses resultados mostram que pelo menos um teste foi encontrado no spark_test.py arquivo, e um ponto (.) significa que um único teste foi encontrado e aprovado. (Um teste com reprovação mostraria um 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)