def test_clean_python_code(repo):
    pack = repo.create_pack('PackName')
    integration = pack.create_integration('integration', 'bla', INTEGRATION_YAML)
    unifier = IntegrationScriptUnifier(str(integration.path))
    script_code = "import demistomock as demisto\nfrom CommonServerPython import *  # test comment being removed\n" \
                  "from CommonServerUserPython import *\nfrom __future__ import print_function"
    # Test remove_print_future is False
    script_code = unifier.clean_python_code(script_code, remove_print_future=False)
    assert script_code == "\n\n\nfrom __future__ import print_function"
    # Test remove_print_future is True
    script_code = unifier.clean_python_code(script_code)
    assert script_code.strip() == ""
def test_insert_script_to_yml(package_path, dir_name, file_path):
    with patch.object(IntegrationScriptUnifier, "__init__", lambda a, b, c, d, e: None):
        unifier = IntegrationScriptUnifier("", None, None, None)
        unifier.package_path = package_path
        unifier.dir_name = dir_name
        unifier.is_script_package = dir_name == 'Scripts'
        with open(file_path + ".yml", "r") as yml:
            test_yml_data = yaml.load(yml)

        test_yml_unified = copy.deepcopy(test_yml_data)

        yml_unified, script_path = unifier.insert_script_to_yml(".py", test_yml_unified, test_yml_data)

        with open(file_path + ".py", mode="r", encoding="utf-8") as script_file:
            script_code = script_file.read()
        clean_code = unifier.clean_python_code(script_code)

        if isinstance(test_yml_unified.get('script', {}), str):
            test_yml_unified['script'] = clean_code
        else:
            test_yml_unified['script']['script'] = clean_code

        assert yml_unified == test_yml_unified
        assert script_path == file_path + ".py"