Compartilhar via

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

  • Artigo

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

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

Para executar pytest no código em um workspace 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 os testes a serem executados. Este exemplo pressupõe que esse arquivo esteja nomeado como spark_test.py e esteja na raiz do projeto do Visual Studio Code. O arquivo contém um pytest acessório, o que torna o SparkSession do cluster (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 esse 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 executor do pytest

Adicione um arquivo Python com o código a seguir, que instrui o pytest a executar os testes da etapa anterior. Este exemplo pressupõe que o arquivo esteja nomeado como pytest_databricks.py e esteja na raiz do projeto do 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 o pytest a executar os testes, você deve criar uma configuração de execução personalizada. Use a configuração de execução baseada em cluster do 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 adicionará um arquivo .vscode/launch.json 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 dessa configuração de execução de Run on Databricks para algum nome de exibição exclusivo para essa configuração; neste exemplo Unit Tests (on Databricks).
    • Altere program de ${file} para o caminho no projeto que contém o executor de teste; neste exemplo ${workspaceFolder}/pytest_databricks.py.
    • Altere args de [] para o caminho no projeto que contém os arquivos com os testes; neste exemplo ["."].

    O arquivo launch.json deve ter esta aparência:

    {
  // 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

Primeiro verifique se o pytest já está instalado no cluster. Por exemplo, com a página de configurações do cluster aberta no workspace do Azure Databricks, faça o seguinte:

  1. Na guia Bibliotecas, se o pytest estiver visível, então o pytest já estará instalado. Se o pytest não estiver visível, clique em Instalar novo.
  2. Em Origem da Biblioteca, clique em PyPI.
  3. Para Pacote, insira pytest.
  4. Clique em Instalar.
  5. Aguarde até que o Status mude de Pendente para Instalado.

Para executar os testes, faça o seguinte em seu projeto do 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 resultados do pytest 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 arquivo spark_test.py e um ponto (.) significa que um único teste foi encontrado e aprovado. (Um teste com falha apareceria com 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)