#!/usr/bin/env python3
"""
测试多主键支持的简单脚本
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from custom_stages.dms_crud_scenario_stage import DmsCrudScenarioStage

def test_build_list_filter_payload():
    """测试LIST过滤条件构建"""
    # 创建一个简单的stage实例来测试方法
    stage = DmsCrudScenarioStage(api_group_metadata={}, apis_in_group=[])
    
    # 测试多主键场景
    identity_id_list = ["wellId", "wellboreId", "eventName"]
    all_pk_values = {
        "wellId": "WELLTL100017525",
        "wellboreId": "WEBHTL100001283", 
        "eventName": "测试井2"
    }
    
    result = stage._build_list_filter_payload(identity_id_list, all_pk_values)
    
    print("=== 多主键LIST过滤条件 ===")
    import json
    print(json.dumps(result, indent=2, ensure_ascii=False))
    
    # 验证结构
    assert "isSearchCount" in result
    assert "query" in result
    assert "filter" in result["query"]
    assert "subFilter" in result["query"]["filter"]
    assert len(result["query"]["filter"]["subFilter"]) == 1  # 简化模式：只使用一个过滤条件
    
    # 验证过滤条件（简化模式：只使用第一个主键）
    assert len(result["query"]["filter"]["subFilter"]) == 1
    sub_filter = result["query"]["filter"]["subFilter"][0]
    assert sub_filter["key"] == identity_id_list[0]  # 第一个主键
    assert sub_filter["symbol"] == "="
    assert sub_filter["realValue"] == [all_pk_values[identity_id_list[0]]]
    
    print("✅ 多主键LIST过滤条件测试通过")

def test_single_pk_scenario():
    """测试单主键场景"""
    stage = DmsCrudScenarioStage(api_group_metadata={}, apis_in_group=[])
    
    # 测试单主键场景
    identity_id_list = ["id"]
    all_pk_values = {"id": "12345"}
    
    result = stage._build_list_filter_payload(identity_id_list, all_pk_values)
    
    print("\n=== 单主键LIST过滤条件 ===")
    import json
    print(json.dumps(result, indent=2, ensure_ascii=False))
    
    # 验证结构
    assert len(result["query"]["filter"]["subFilter"]) == 1
    assert result["query"]["filter"]["subFilter"][0]["key"] == "id"
    assert result["query"]["filter"]["subFilter"][0]["realValue"] == ["12345"]
    
    print("✅ 单主键LIST过滤条件测试通过")

if __name__ == "__main__":
    test_build_list_filter_payload()
    test_single_pk_scenario()
    print("\n🎉 所有测试通过！")