单元测试中隔离依赖的最佳实践：Databricks Python模块测试难题

先澄清你的误解：方案2和3其实能实现依赖隔离 #

你觉得方案2、3会保留Spark依赖导致失败，其实是没找对实现方式。这两个方案完全可以做到在单元测试时不触发Spark runtime的连接尝试，核心是延迟导入或者mock掉整个databricks.sdk.runtime模块的导入。

方案1：拆分依赖与非依赖代码（最直观的方式） #

把utils.py拆成两个模块：

• core_utils.py：存放所有不依赖Spark/dbutil的工具函数，比如字符串处理、数据校验逻辑等
• databricks_utils.py：只放需要Spark/dbutil的代码，开头保留from databricks.sdk.runtime import spark, dbutil

单元测试时只导入core_utils.py里的函数，完全避开Spark依赖；在Databricks环境中，业务代码可以同时导入两个模块的内容。

示例结构：

your_project/
├── core_utils.py
├── databricks_utils.py
└── test/
    └── test_core_utils.py

core_utils.py示例：

def clean_string(input_str: str) -> str:
    return input_str.strip().lower()

databricks_utils.py示例：

from databricks.sdk.runtime import spark, dbutil

def write_df_to_delta(df, path: str):
    df.write.format("delta").mode("overwrite").save(path)

测试时直接导入core_utils.clean_string，完全不会触发Spark相关的导入。

方案2：策略模式+延迟导入（兼顾环境适配） #

核心是把Spark相关的操作抽象成接口，在不同环境下加载不同实现，并且延迟导入Spark依赖，只有在Databricks环境执行时才会触发导入。

步骤：

1. 定义抽象操作接口
2. 分别实现Databricks版本和测试版本的逻辑
3. 根据环境变量判断加载哪个实现，且只在需要时导入Spark相关代码

示例代码：
utils.py：

import os

class DataWriter:
    def write(self, data, path):
        raise NotImplementedError

class DeltaDataWriter(DataWriter):
    def __init__(self):
        # 延迟导入，只有实例化这个类时才会加载Spark
        from databricks.sdk.runtime import spark
        self.spark = spark

    def write(self, data, path):
        df = self.spark.createDataFrame(data)
        df.write.format("delta").mode("overwrite").save(path)

class MockDataWriter(DataWriter):
    def write(self, data, path):
        # 测试时的模拟实现，比如把数据写入本地文件或内存
        with open(path, "w") as f:
            f.write(str(data))

def get_data_writer() -> DataWriter:
    if os.getenv("ENV") == "DATABRICKS":
        return DeltaDataWriter()
    else:
        return MockDataWriter()

单元测试时，设置ENV=TEST（或默认走Mock），此时只会加载MockDataWriter，完全不会触发databricks.sdk.runtime的导入，也就不会尝试连接Spark。

方案3：Mock整个Spark依赖（适合不想改太多原代码的场景） #

用unittest.mock的patch来mock掉databricks.sdk.runtime模块，让导入时不会真正去加载Spark runtime。关键是在导入utils模块之前就完成mock，因为Python的导入是一次性的，一旦开始导入utils，再mock就晚了。

示例test.py：

import unittest
from unittest.mock import Mock, patch

# 在导入utils之前mock databricks.sdk.runtime
with patch('databricks.sdk.runtime') as mock_runtime:
    # mock spark和dbutil对象
    mock_spark = Mock()
    mock_dbutil = Mock()
    mock_runtime.spark = mock_spark
    mock_runtime.dbutil = mock_dbutil

    # 现在可以安全导入utils了
    from utils import your_spark_dependent_function

class TestUtils(unittest.TestCase):
    def test_your_function(self):
        # 模拟spark相关方法的返回值
        mock_spark.createDataFrame.return_value = Mock()
        # 执行测试
        your_spark_dependent_function()
        # 验证调用
        mock_spark.createDataFrame.assert_called_once()

if __name__ == '__main__':
    unittest.main()

这样做的话，导入databricks.sdk.runtime时，实际拿到的是你mock的对象，不会触发真实的Spark连接逻辑。

总结 #

• 如果代码拆分成本低，优先选方案1，最清晰也最容易维护
• 如果需要同一套接口适配多环境，选方案2，扩展性更好
• 如果不想改动原有utils代码，选方案3，快速实现测试隔离

内容的提问来源于stack exchange，提问作者absmiths