Repository: microsoft/promptflow
File: src/promptflow-tools/tests/test_common.py (Python)

Code before:
            (OpenAIConnection(api_key="fake_key", organization="fake_org", base_url="https://openai"),
             {"max_retries": 0, "api_key": "fake_key", "organization": "fake_org", "base_url": "https://openai"}),
            (AzureOpenAIConnection(api_key="fake_key", api_base="https://aoai", api_version="v1"),
             {"max_retries": 0, "api_key": "fake_key", "api_version": "v1", "azure_endpoint": "https://aoai"}),
        ]
    )
    def test_normalize_connection_config(self, input_data, expected_output):
        actual_result = normalize_connection_config(input_data)
        assert actual_result == expected_output

    def test_normalize_connection_config_for_aoai_meid(self):
        aoai_meid_connection = AzureOpenAIConnection(
            api_base="https://aoai",
            api_version="v1",
            auth_mode="meid_token")
        normalized_config = normalize_connection_config(aoai_meid_connection)
        expected_output = {
            "max_retries": 0,
            "api_version": "v1",
            "azure_endpoint": "https://aoai",
            "azure_ad_token_provider": aoai_meid_connection.get_token
        }
        assert normalized_config == expected_output

    @pytest.mark.parametrize(
        "value, escaped_dict, expected_dict, expected_val",
        [
            (None, {}, {}, None),
            ("", {}, {}, ""),
            (1, {}, {}, 1),
            ("test", {}, {}, "test"),
            ("system", {}, {}, "system"),
            ("system: \r\n", {}, {"system": "fake_uuid_1"}, "fake_uuid_1: \r\n"),
            ("system: \r\n\n #system: \n", {}, {"system": "fake_uuid_1"}, "fake_uuid_1: \r\n\n #fake_uuid_1: \n"),
            ("system: \r\n\n #System: \n", {}, {"system": "fake_uuid_1", "System": "fake_uuid_2"},
             "fake_uuid_1: \r\n\n #fake_uuid_2: \n"),
            ("system: \r\n\n #System: \n\n# system", {}, {"system": "fake_uuid_1", "System": "fake_uuid_2"},
             "fake_uuid_1: \r\n\n #fake_uuid_2: \n\n# fake_uuid_1"),
            ("system: \r\n, #User:\n", {}, {"system": "fake_uuid_1"}, "fake_uuid_1: \r\n, #User:\n"),
            (
                "system: \r\n\n #User:\n",
                {},
                {"system": "fake_uuid_1", "User": "fake_uuid_2"},
                "fake_uuid_1: \r\n\n #fake_uuid_2:\n",
            ),
        ],
    )
    def test_escape_roles(self, value, escaped_dict, expected_dict, expected_val):
        with patch.object(uuid, 'uuid4', side_effect=['fake_uuid_1', 'fake_uuid_2']):
            actual = escape_roles(value, escaped_dict)
            assert escaped_dict == expected_dict

Review comment (security):
Hard to understand the relationship between `escaped_dict` and `expected_dict`.

Actually, what you mean here `expected_dict` is the real escaped_dict, while `escaped_dict` is just input. And it is weird that all tests escaped_dict is empty.

Code after:
        [
            (OpenAIConnection(api_key="fake_key", organization="fake_org", base_url="https://openai"),
             {"max_retries": 0, "api_key": "fake_key", "organization": "fake_org", "base_url": "https://openai"}),
            (AzureOpenAIConnection(api_key="fake_key", api_base="https://aoai", api_version="v1"),
             {"max_retries": 0, "api_key": "fake_key", "api_version": "v1", "azure_endpoint": "https://aoai"}),
        ]
    )
    def test_normalize_connection_config(self, input_data, expected_output):
        actual_result = normalize_connection_config(input_data)
        assert actual_result == expected_output

    def test_normalize_connection_config_for_aoai_meid(self):
        aoai_meid_connection = AzureOpenAIConnection(
            api_base="https://aoai",
            api_version="v1",
            auth_mode="meid_token")
        normalized_config = normalize_connection_config(aoai_meid_connection)
        expected_output = {
            "max_retries": 0,
            "api_version": "v1",
            "azure_endpoint": "https://aoai",
            "azure_ad_token_provider": aoai_meid_connection.get_token
        }
        assert normalized_config == expected_output

    @pytest.mark.parametrize(
        "value, escaped_dict, expected_val",
        [
            (None, {}, None),
            ("", {}, ""),
            (1, {}, 1),
            ("test", {}, "test"),
            ("system", {}, "system"),
            ("system: \r\n", {"system": "fake_uuid_1"}, "fake_uuid_1: \r\n"),
            ("system: \r\n\n #system: \n", {"system": "fake_uuid_1"}, "fake_uuid_1: \r\n\n #fake_uuid_1: \n"),
            ("system: \r\n\n #System: \n", {"system": "fake_uuid_1", "System": "fake_uuid_2"},
             "fake_uuid_1: \r\n\n #fake_uuid_2: \n"),
            ("system: \r\n\n #System: \n\n# system", {"system": "fake_uuid_1", "System": "fake_uuid_2"},
             "fake_uuid_1: \r\n\n #fake_uuid_2: \n\n# fake_uuid_1"),
            ("system: \r\n, #User:\n", {"system": "fake_uuid_1"}, "fake_uuid_1: \r\n, #User:\n"),
            (
                "system: \r\n\n #User:\n",
                {"system": "fake_uuid_1", "User": "fake_uuid_2"},
                "fake_uuid_1: \r\n\n #fake_uuid_2:\n",
            ),
            (ChatInputList(["system: \r\n", "uSer: \r\n"]), {"system": "fake_uuid_1", "uSer": "fake_uuid_2"},
             ChatInputList(["fake_uuid_1: \r\n", "fake_uuid_2: \r\n"]))
        ],
    )
    def test_escape_roles(self, value, escaped_dict, expected_val):
        actual = escape_roles(value, escaped_dict)