如何在Pandas测试中为枚举列构造正确的模拟数据？

我在transformer.py中处理包含id和createdAt列的DataFrame，通过以下代码新增第三列：

class ColorType(Enum):
      LIGHT = 'LIGHT'
      def __str__(self):
          return self.name
data["colorType"] = ColorType

为测试该转换器，我在test_transformer.py中编写了如下代码（注：原代码存在语法错误，已修正括号）：

record = utils.EXTRACTED_RECORD # 模拟原始数据
result = transformer.transform_columns(record)
expected = pd.DataFrame([
    {"id": 111111, 
     "completedAt":"2030-06-26T19:08:38.056Z",
     "colorType":"(LIGHT)"}
])
testing.assert_frame_equal(result, expected)

运行测试时出现错误：

AssertionError: DataFrame.iloc[:, 2] (column name="colorType" values are different (100.0%)
[index]: [0, 1]
[left]: [(LIGHT), (LIGHT)]
[right]: [(LIGHT), (LIGHT)]

我需要知道如何在预期DataFrame中正确模拟colorType枚举列，让断言通过？移除该列时测试可正常运行。

核心问题是实际生成的列类型和预期值类型不匹配，同时原transformer代码的枚举赋值逻辑有误。

1. 修正transformer中的枚举赋值逻辑 #

你当前代码data["colorType"] = ColorType是把整个枚举类赋值给每一行，这不是正确用法。应该赋值枚举实例ColorType.LIGHT：

from enum import Enum
import pandas as pd

class ColorType(Enum):
      LIGHT = 'LIGHT'
      def __str__(self):
          return self.name

# 正确赋值枚举实例
data["colorType"] = ColorType.LIGHT

2. 统一测试用例的预期值类型 #

根据实际生成的列类型，调整预期值：

场景A：预期列存储枚举实例

如果transformer中直接存储枚举实例，测试里的预期值也要用ColorType.LIGHT：

import pandas as pd
from transformer import ColorType

record = utils.EXTRACTED_RECORD
result = transformer.transform_columns(record)
expected = pd.DataFrame([
    {"id": 111111, 
     "completedAt":"2030-06-26T19:08:38.056Z",
     "colorType": ColorType.LIGHT}
])
testing.assert_frame_equal(result, expected)

场景B：预期列存储字符串

如果需要把枚举转成字符串存储（比如后续要序列化），先修改transformer代码：

# transformer.py中调整
data["colorType"] = str(ColorType.LIGHT)
# 或者用枚举的value属性：data["colorType"] = ColorType.LIGHT.value

然后测试用例的预期值改为对应字符串：

expected = pd.DataFrame([
    {"id": 111111, 
     "completedAt":"2030-06-26T19:08:38.056Z",
     "colorType": "LIGHT"}
])
testing.assert_frame_equal(result, expected)

补充说明 #

错误提示里两边显示的(LIGHT)只是打印格式相同，实际左边是ColorType类对象，右边是字符串，类型不匹配才导致断言失败。必须保证预期值的类型和实际列的类型完全一致。

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