测试 Databricks ODBC 驱动程序
本文介绍如何测试使用 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
。