Compartir a través de

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

En esta página se describe cómo ejecutar pruebas de Python mediante la extensión Databricks para Visual Studio Code. Consulte ¿Qué es la extensión de Databricks para Visual Studio Code?.

Ejecución de pruebas mediante pytest

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 pytestaccesorio, que hace que el clúster SparkSession (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:

  1. En el menú principal, haga clic en Ejecutar > Agregar configuración.

  2. 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.

  3. 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:

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

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

  1. En el menú principal, haga clic en Ver > Ejecutar.
  2. En la lista Ejecutar y depurar, haga clic en Pruebas unitarias (en Databricks) si aún no está seleccionada.
  3. 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)

Ejecución de pruebas mediante Databricks Connect

Para ejecutar pruebas localmente que usan las API de Spark, use Databricks Connect.

Paso 1: Configuración de Databricks Connect

Siga los pasos para configurar Databricks Connect para la extensión. Consulte Depurar código mediante Databricks Connect en la extensión de Databricks para Visual Studio Code.

Paso 2: Crear una prueba unitaria

Agregue un archivo de Python con el código siguiente, que contiene la prueba que se va a ejecutar. En este ejemplo se supone que este archivo se denomina main_test.py.

from my_project import main


def test_find_all_taxis():
    taxis = main.find_all_taxis()
    assert taxis.count() > 5

Paso 3: Agregar o actualizar la configuración de inicio debugpy

A continuación, cree una debugpy configuración de inicio que habilite Databricks Connect.

  1. En el menú principal de Visual Studio Code, haga clic en Ejecutar > agregar configuración.

  2. En la paleta de comandos, seleccione Depurador de Python.

    Visual Studio Code agrega un archivo .vscode/launch.json al proyecto, si este archivo aún no existe.

  3. Agregar "databricks": true campo. Esto habilita Databricks Connect.

{
  "version": "0.2.0",
  "configurations": [
    {
      "name": "Unit Tests (on Databricks)",
      "type": "debugpy",
      "databricks": true,
      "request": "launch",
      "program": "${file}",
      "args": ["."],
      "env": {},
      "console": "integratedTerminal"
    }
  ]
}

Paso 4: Ejecución de las pruebas

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

  1. En el menú principal, haga clic en Ver > pruebas para abrir el panel de pruebas.
  2. En el panel de pruebas, ejecute la prueba haciendo clic en el icono de depuración asociado a main_test.py. Tenga en cuenta que al ejecutar la prueba no se desencadenará la configuración de depuración modificada y el código no tendrá acceso a Databricks Connect.
  • Last updated on