Ejecución de pruebas con pytest para la extensión de Databricks para Visual Studio Code

Artículo
03/01/2024

En este artículo, se describe cómo ejecutar pruebas mediante pytest para la extensión Databricks para Visual Studio Code. Consulte ¿Qué es la extensión de Databricks para Visual Studio Code?.

En esta información, se asume que ya ha instalado y configurado la extensión de Databricks para Visual Studio Code. Consulte Instalación de la extensión de Databricks para Visual Studio Code.

Puede ejecutar pytest en código local que no necesita una conexión a un clúster en un área de trabajo remota de Azure Databricks. Por ejemplo, puede usar pytest para probar las funciones que aceptan y devuelven dataframes de PySpark en la memoria local. Para empezar a trabajar con pytest y ejecutarlo localmente, consulte Introducción en la documentación de pytest.

Para ejecutar pytest en código en un área de trabajo remota de Azure Databricks, haga lo siguiente en el proyecto de Visual Studio Code:

Paso 1: Creación de las pruebas

Agregue un archivo de Python con el código siguiente, que contiene las pruebas que se van a ejecutar. En este ejemplo se asume que este archivo se denomina spark_test.py y está en la raíz del proyecto de Visual Studio Code. Este archivo contiene un accesorio pytest, lo que hace que el SparkSession del clúster (el punto de entrada a la funcionalidad de Spark en el clúster) esté disponible para las pruebas. Este archivo contiene una única prueba que comprueba si la celda especificada de la tabla contiene el valor especificado. Puede agregar sus propias pruebas a este archivo según sea necesario.

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'

Paso 2: Creación del ejecutor pytest

Agregue un archivo de Python con el código siguiente, que indica a pytest que ejecute las pruebas del paso anterior. En este ejemplo se asume que el archivo se denomina pytest_databricks.py y está en la raíz del proyecto de 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:])

Paso 3: Creación de una configuración de ejecución personalizada

Para indicar a pytest que ejecute las pruebas, debe crear una configuración de ejecución personalizada. Use la configuración de ejecución basada en clústeres de Databricks existente para crear su propia configuración de ejecución personalizada, como se indica a continuación:

En el menú principal, haga clic en Ejecutar > Agregar configuración.
En la Paleta de comandos, seleccione Databricks.

Visual Studio Code agrega un archivo .vscode/launch.json al proyecto, si este archivo aún no existe.
Cambie la configuración de ejecución de inicio como se indica a continuación y guarde el archivo:
- Cambie el nombre de esta configuración de ejecución de Run on Databricks a algún nombre para mostrar único para esta configuración, en este ejemplo Unit Tests (on Databricks).
- Cambie program de ${file} a la ruta de acceso del proyecto que contiene el ejecutor de pruebas, en este ejemplo ${workspaceFolder}/pytest_databricks.py.
- Cambie args de [] a la ruta de acceso del proyecto que contiene los archivos con las pruebas, en este ejemplo ["."].
El archivo launch.json debería tener este aspecto:
```
{
  // 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": {}
    }
  ]
}
```

Paso 4: Ejecución de las pruebas

Asegúrese de que pytest ya está instalado en el clúster en primer lugar. Por ejemplo, con la página de configuración del clúster abierta en el área de trabajo de Azure Databricks, haga lo siguiente:

En la pestaña Bibliotecas, si pytest está visible, pytest ya está instalado. Si pytest no está visible, haga clic en Instalar nuevo.
En Origen de biblioteca, haga clic en PyPI.
En Paquete, escriba pytest.
Haga clic en Instalar.
Espere hasta que el Estado cambie de Pendiente a Instalado.

Para ejecutar las pruebas, haga lo siguiente desde el proyecto de Visual Studio Code:

En el menú principal, haga clic en Ver > Ejecutar.
En la lista Ejecutar y depurar, haga clic en Pruebas unitarias (en Databricks) si aún no está seleccionada.
Haga clic en el icono de flecha verde (Iniciar depuración).

Los resultados pytest se muestran en la Consola de depuración (Ver > Consola de depuración en el menú principal). Por ejemplo, estos resultados muestran que se encontró al menos una prueba en el archivo spark_test.py, y un punto (.) significa que se encontró y se superó una sola prueba. (Una prueba con errores mostraría un 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)

Compartir a través de