为使用DynamicFrame的Glue代码编写单元测试的相关问题咨询

快速生成DynamicFrame #

你可以直接基于已有的Spark DataFrame转换得到DynamicFrame，完全复用你创建Spark DF的逻辑，不需要单独构造DynamicFrame：

from awsglue.dynamicframe import DynamicFrame

# 其中glue_context为GlueContext实例，第三个参数为DynamicFrame的名称可自定义
test_dyf = DynamicFrame.fromDF(your_spark_df, glue_context, "test_dynamic_frame")

初始化/Mock GlueContext的方案 #

实际单元测试不需要做mock，本地PySpark环境可以直接实例化原生的GlueContext运行，不需要连接AWS服务，测试结果和线上环境完全一致：

1. 首先安装本地依赖：pip install aws-glue-libs（也可直接使用AWS官方提供的Glue本地Docker镜像）
2. 在测试类的初始化方法中生成GlueContext实例即可，示例代码如下：

import unittest
from pyspark.sql import SparkSession
from awsglue.context import GlueContext

class TestGlueDBReader(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # 初始化本地Spark会话
        cls.spark = SparkSession.builder \
            .master("local[*]") \
            .appName("GlueUnitTest") \
            .getOrCreate()
        # 直接实例化GlueContext，无网络请求
        cls.glue_context = GlueContext(cls.spark.sparkContext)
        # 初始化待测的业务对象，传入GlueContext
        cls.gluedbreader = 你的业务类(glue_context=cls.glue_context)
        # 提前初始化测试用基础DataFrame
        cls.df = cls.spark.createDataFrame([], schema="你的业务Schema")

如果有特殊场景必须做Mock，可以用unittest.mock.patch模拟你用到的GlueContext方法，自定义返回值即可。

DynamicFrame断言校验方法 #

你之前直接对比DynamicFrame和Spark DF对象失败，是因为二者属于不同的类实例，正确的校验逻辑是先把DynamicFrame转成Spark DF再做内容和Schema的对比：

基础断言方法（无第三方依赖） #

你可以封装一个公共的断言方法复用：

from pyspark.sql import DataFrame
from awsglue.dynamicframe import DynamicFrame

def assert_dyf_matches_df(self, dyf: DynamicFrame, expected_df: DataFrame):
    # 统一排序，避免行顺序影响对比结果
    actual_sorted = dyf.toDF().orderBy(*expected_df.columns)
    expected_sorted = expected_df.orderBy(*expected_df.columns)
    # 先校验Schema一致性
    self.assertEqual(actual_sorted.schema, expected_sorted.schema)
    # 再校验数据内容一致性
    self.assertEqual(actual_sorted.collect(), expected_sorted.collect())

修改后的测试用例如下：

def test_set_dynamic_frame(self):
    newRow = self.spark.createDataFrame([('test2', '24/12/1991', 'abcd', '24/12/1991', 'ábcd', '24/12/1991', '1')], schema=self.df.schema)
    df = self.df.union(newRow)

    self.gluedbreader.set_dynamic_frame(df)
    self.assert_dyf_matches_df(self.gluedbreader.get_dynamic_frame, df)

简化断言方案（使用第三方测试库） #

可以使用专门针对Spark单元测试的chispa库简化对比逻辑，自动处理空值、行顺序、浮点数精度等特殊场景：

from chispa import assert_df_equality

def test_set_dynamic_frame(self):
    newRow = self.spark.createDataFrame([('test2', '24/12/1991', 'abcd', '24/12/1991', 'ábcd', '24/12/1991', '1')], schema=self.df.schema)
    df = self.df.union(newRow)

    self.gluedbreader.set_dynamic_frame(df)
    actual_df = self.gluedbreader.get_dynamic_frame.toDF()
    assert_df_equality(actual_df, df, ignore_row_order=True)

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