测试 Databricks ODBC 驱动程序

项目
06/24/2024

本文介绍如何测试使用 Databricks ODBC 驱动程序的代码。

要测试使用 Databricks ODBC 驱动程序和 DSN 或无 DSN 连接字符串的代码，可以使用支持 ODBC 的编程语言的常用测试框架。例如，以下示例 Python 代码使用 pyodbc、pytest 和 unittest.mock，通过 DSN 自动执行和测试 Databricks ODBC 驱动程序。此示例代码基于将 Python 和 pyodbc 连接到 Azure Databricks 中的示例代码。

以下名为 helpers.py 的示例代码文件包含几个使用 DSN 自动执行 Databricks 驱动程序的函数：

connect_to_dsn 函数使用 DSN 通过 Azure Databricks 计算资源打开连接。
get_cursor_from_connection 函数使用连接获取游标，从而可以通过计算资源对数据进行提取操作。
select_from_nyctaxi_trips 函数使用游标从 samples 目录 nyctaxi 架构中的 trips 表中选择指定数量的数据行。
print_rows 函数将数据行的内容输出到屏幕上。

# helpers.py

from pyodbc import connect, Connection, Cursor

def connect_to_dsn(
  connstring: str,
  autocommit: bool
) -> Connection:

  connection = connect(
    connstring,
    autocommit = autocommit
  )

  return connection

def get_cursor_from_connection(
  connection: Connection
) -> Cursor:

  cursor = connection.cursor()
  return cursor

def select_from_nyctaxi_trips(
  cursor: Cursor,
  num_rows: int
) -> Cursor:

  select_cursor = cursor.execute(f"SELECT * FROM samples.nyctaxi.trips LIMIT {num_rows}")
  return select_cursor

def print_rows(cursor: Cursor):
  for row in cursor.fetchall():
    print(row)

下面名为 main.py 的示例代码文件调用 helpers.py 文件中的函数：

# main.py

from helpers import *

connection = connect_to_dsn(
  connstring = "DSN=<your-dsn-name>",
  autocommit = True
)

cursor = get_cursor_from_connection(
  connection = connection)

select_cursor = select_from_nyctaxi_trips(
  cursor = cursor,
  num_rows = 2
)

print_rows(
  cursor = select_cursor
)

下面名为 test_helpers.py 的示例代码文件使用 pytest 测试 helpers.py 文件中的函数。下面的示例代码没有使用实际计算资源的时间和成本来调用 helpers.py 文件中的函数，而是使用 unittest.mock 来模拟这些调用。这些模拟调用通常在几秒钟内完成，增加了你对代码质量的信心，同时不会更改现有 Azure Databricks 帐户或工作区的状态。

# test_helpers.py

from pyodbc import SQL_DBMS_NAME
from helpers import *
from unittest.mock import patch
import datetime

@patch("helpers.connect_to_dsn")
def test_connect_to_dsn(mock_connection):
  mock_connection.return_value.getinfo.return_value = "Spark SQL"

  mock_connection = connect_to_dsn(
    connstring = "DSN=<your-dsn-name>",
    autocommit = True
  )

  assert mock_connection.getinfo(SQL_DBMS_NAME) == "Spark SQL"

@patch('helpers.get_cursor_from_connection')
def test_get_cursor_from_connection(mock_connection):
  mock_cursor = mock_connection.return_value.cursor
  mock_cursor.return_value.rowcount = -1

  mock_connection = connect_to_dsn(
    connstring = "DSN=<your-dsn-name>",
    autocommit = True
  )

  mock_cursor = get_cursor_from_connection(
    connection = mock_connection
  )

  assert mock_cursor.rowcount == -1

@patch('helpers.select_from_nyctaxi_trips')
def test_select_from_nyctaxi_trips(mock_connection):
  mock_cursor = mock_connection.return_value.cursor
  mock_get_cursor = mock_cursor.return_value.execute
  mock_select_cursor = mock_get_cursor.return_value.arraysize = 1

  mock_connection = connect_to_dsn(
    connstring = "DSN=<your-dsn-name>",
    autocommit = True
  )

  mock_get_cursor = get_cursor_from_connection(
    connection = mock_connection
  )

  mock_select_cursor = select_from_nyctaxi_trips(
    cursor = mock_get_cursor,
    num_rows = 2
  )

  assert mock_select_cursor.arraysize == 1

@patch('helpers.print_rows')
def test_print_rows(mock_connection, capsys):
  mock_cursor = mock_connection.return_value.cursor
  mock_get_cursor = mock_cursor.return_value.execute
  mock_select_cursor = mock_get_cursor.return_value.fetchall.return_value = [
    (datetime.datetime(2016, 2, 14, 16, 52, 13), datetime.datetime(2016, 2, 14, 17, 16, 4), 4.94, 19.0, 10282, 10171),
    (datetime.datetime(2016, 2, 4, 18, 44, 19), datetime.datetime(2016, 2, 4, 18, 46), 0.28, 3.5, 10110, 10110)
  ]

  mock_connection = connect_to_dsn(
    connstring = "DSN=<your-dsn-name>",
    autocommit = True
  )

  mock_get_cursor = get_cursor_from_connection(
    connection = mock_connection
  )

  mock_select_cursor = select_from_nyctaxi_trips(
    cursor = mock_get_cursor,
    num_rows = 2
  )

  print_rows(
    cursor = mock_select_cursor
  )

  captured = capsys.readouterr()
  assert captured.out == "(datetime.datetime(2016, 2, 14, 16, 52, 13), datetime.datetime(2016, 2, 14, 17, 16, 4), 4.94, 19.0, 10282, 10171)\n" \
                         "(datetime.datetime(2016, 2, 4, 18, 44, 19), datetime.datetime(2016, 2, 4, 18, 46), 0.28, 3.5, 10110, 10110)\n"

select_from_nyctaxi_trips 函数包含 SELECT 语句，因此不会更改 trips 表的状态，在此示例中并不是一定需要模拟。但是，模拟让你能够快速运行测试，无需等待与计算资源建立实际连接。此外，通过模拟，可以多次针对可能更改表状态的函数运行模拟测试，例如 INSERT INTO、UPDATE 和 DELETE FROM。

通过

测试 Databricks ODBC 驱动程序

其他资源