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

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

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

Para executar pytest no código em um site remoto Databricks workspace, faça o seguinte em seu projeto do Visual Studio Code:

passo 1: Criar os testes

Adicione um arquivo Python com o seguinte código, que contém seus testes para execução. Este exemplo pressupõe que esse arquivo tenha o nome spark_test.py e esteja na raiz do seu projeto do Visual Studio Code. Esse arquivo contém um fixture pytest, 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'

passo 2: criar o executor pytest

Adicione um arquivo Python com o código a seguir, que instrui pytest a executar seus testes da passo anterior. Este exemplo supõe que o arquivo tenha o nome pytest_databricks.py e esteja na raiz do seu 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:])

passo 3: criar uma configuração de execução personalizada

Para instruir pytest a executar seus testes, o senhor 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, como segue:

No menu principal, clique em execução > Add configuration.
Na paleta de comandos, selecione Databricks.

O Visual Studio Code adiciona um arquivo .vscode/launch.json ao seu projeto, se esse arquivo ainda não existir.
Altere a configuração de execução do starter da seguinte forma 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 seus testes, neste exemplo ["."].
Seu arquivo launch.json deve ter a seguinte 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": {}
    }
  ]
}
```

passo 4: execução dos testes

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

Na biblioteca tab, se o pytest estiver visível, então pytest já está instalado. Se o pytest não estiver visível, clique em Instalar novo.
Para obter o código-fonte da biblioteca, clique em PyPI.
Para pacote, digite pytest.
Clique em Instalar.
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:

No menu principal, clique em view > execução.
Na lista Execução e depuração, clique em Testes de unidade (em Databricks), se ainda não estiver selecionado.
Clique no íconeda seta verde(começar depuração).

Os resultados pytest são exibidos no Console de depuração(veja > Debug Console 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 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)