PROCESS_TREE_RAW = [
    {
        "id": 3,
        "ptid": 3,
        "pid": 1,
        "name": "1: <Pruned Process>",
        "parent": "4: System",
        "children": [
            {
                "id": 44,
                "ptid": 44,
                "pid": 4236,
                "name": "4236: mmc.exe",
                "parent": "1: <Pruned Process>",
                "children": []
            },
            {
                "id": 45,
                "ptid": 45,
                "pid": 4840,
                "name": "4840: cmd.exe",
                "parent": "1: <Pruned Process>",
                "children": []
            }
        ]
    }
]

PROCESS_TREE_TWO_GENERATIONS_RAW = [
    {
        "id": 3,
        "ptid": 3,
        "pid": 1,
        "name": "1: <Pruned Process>",
        "parent": "4: System",
        "children": [
            {
                "id": 44,
                "ptid": 44,
                "pid": 4236,
                "name": "4236: mmc.exe",
                "parent": "1: <Pruned Process>",
                "children": [
                    {
                        "id": 420,
                        "ptid": 44,
                        "pid": 4236,
                        "name": "4236: mmc.exe",
                        "parent": "1: <Pruned Process>",
                        "children": []
                    }
                ]
            }
        ]
    }
]

PROCESS_TREE_ITEM_RES = {
    "ID": 3,
    "PTID": 3,
    "PID": 1,
    "Name": "1: <Pruned Process>",
    "Parent": "4: System",
    "Children": [
        {
            "ID": 44,
            "PTID": 44,
            "PID": 4236,
            "Name": "4236: mmc.exe",
            "Parent": "1: <Pruned Process>",
            "Children": []
        },
        {
            "ID": 45,
            "PTID": 45,
            "PID": 4840,
            "Name": "4840: cmd.exe",
            "Parent": "1: <Pruned Process>",
            "Children": []
        }
    ]
}

PROCESS_TREE_ITEM_TWO_GENERATIONS_RES = {
    "ID": 3,
    "PTID": 3,
    "PID": 1,
    "Name": "1: <Pruned Process>",
    "Parent": "4: System",
    "Children": [
        {
            "ID": 44,
            "PTID": 44,
            "PID": 4236,
            "Name": "4236: mmc.exe",
            "Parent": "1: <Pruned Process>",
            "Children": [
                {
                    "id": 420,
                    "ptid": 44,
                    "pid": 4236,
                    "name": "4236: mmc.exe",
                    "parent": "1: <Pruned Process>",
                    "children": []
                }
            ]
        }
    ]
}

PROCESS_TREE_READABLE_RES = {
    "ID": 3,
    "PTID": 3,
    "PID": 1,
    "Name": "1: <Pruned Process>",
    "Parent": "4: System",
    "Children": [
        {
            "ID": 44,
            "PTID": 44,
            "PID": 4236,
            "Name": "4236: mmc.exe",
            "Parent": "1: <Pruned Process>",
            "ChildrenCount": 0
        },
        {
            "ID": 45,
            "PTID": 45,
            "PID": 4840,
            "Name": "4840: cmd.exe",
            "Parent": "1: <Pruned Process>",
            "ChildrenCount": 0
        }
    ]
}

PROCESS_TREE_TWO_GENERATIONS_READABLE_RES = {
    "ID": 3,
    "PTID": 3,
    "PID": 1,
    "Name": "1: <Pruned Process>",
    "Parent": "4: System",
    "Children": [
        {
            "ID": 44,
            "PTID": 44,
            "PID": 4236,
            "Name": "4236: mmc.exe",
            "Parent": "1: <Pruned Process>",
            "ChildrenCount": 1
        }
    ]
}


def test_get_process_tree_item():
    from TaniumThreatResponse import get_process_tree_item

    tree, readable_output = get_process_tree_item(PROCESS_TREE_RAW[0], 0)

    assert tree == PROCESS_TREE_ITEM_RES
    assert readable_output == PROCESS_TREE_READABLE_RES


def test_get_process_tree_item_two_generations():
    from TaniumThreatResponse import get_process_tree_item

    tree, readable_output = get_process_tree_item(PROCESS_TREE_TWO_GENERATIONS_RAW[0], 0)

    assert tree == PROCESS_TREE_ITEM_TWO_GENERATIONS_RES
    assert readable_output == PROCESS_TREE_TWO_GENERATIONS_READABLE_RES
