Repository: microsoft/promptflow
File: src/promptflow-devkit/tests/sdk_cli_test/e2etests/test_flow_save.py (Python)

Code before:
        target_path = f"{FLOWS_DIR}/saved/hello_callable"
        if os.path.exists(target_path):
            shutil.rmtree(target_path)
        pf.flows._save(
            entry=GlobalHello,
            path=target_path,
        )

        flow = load_flow(target_path)
        assert flow._data == {
            "entry": "test_flow_save:GlobalHello",
            "init": {
                "connection": {
                    "type": "AzureOpenAIConnection",
                }
            },
            "inputs": {
                "text": {
                    "type": "string",
                }
            },
        }

    def test_pf_infer_signature_include_primitive_output(self):
        pf = PFClient()
        flow_meta = pf.flows.infer_signature(entry=global_hello)
        assert flow_meta == {
            "inputs": {
                "text": {
                    "type": "string",
                }
            },
            "outputs": {
                "output": {
                    "type": "string",
                }
            },
        }

    def test_pf_save_callable_function(self):
        pf = PFClient()
        target_path = f"{FLOWS_DIR}/saved/hello_callable"
        if os.path.exists(target_path):
            shutil.rmtree(target_path)
        pf.flows._save(
            entry=global_hello,
            path=target_path,
        )

        flow = load_flow(target_path)
        assert flow._data == {

Review comment (suggestion):
the public API `pf.flows.infer_signature` is for user and shouldn't include primitive output (we don't want customer be aware of this)

Code after:
        target_path = f"{FLOWS_DIR}/saved/hello_callable"
        if os.path.exists(target_path):
            shutil.rmtree(target_path)
        pf.flows._save(
            entry=GlobalHello,
            path=target_path,
        )

        flow = load_flow(target_path)
        assert flow._data == {
            "entry": "test_flow_save:GlobalHello",
            "init": {
                "connection": {
                    "type": "AzureOpenAIConnection",
                }
            },
            "inputs": {
                "text": {
                    "type": "string",
                }
            },
        }

    def test_pf_infer_signature_include_primitive_output(self):
        pf = PFClient()
        flow_meta = pf.flows._infer_signature(entry=global_hello, include_primitive_output=True)
        assert flow_meta == {
            "inputs": {
                "text": {
                    "type": "string",
                }
            },
            "outputs": {
                "output": {
                    "type": "string",
                }
            },
        }

    def test_pf_save_callable_function(self):
        pf = PFClient()
        target_path = f"{FLOWS_DIR}/saved/hello_callable"
        if os.path.exists(target_path):
            shutil.rmtree(target_path)
        pf.flows._save(
            entry=global_hello,
            path=target_path,
        )

        flow = load_flow(target_path)
        assert flow._data == {