test_trace.py

import datetime
import json
import platform
import sys
import time
import typing
import uuid
from pathlib import Path
from unittest.mock import patch

import pytest
from _constants import PROMPTFLOW_ROOT
from mock import mock

from promptflow._constants import (
    RUNNING_LINE_RUN_STATUS,
    SpanAttributeFieldName,
    SpanResourceAttributesFieldName,
    SpanResourceFieldName,
)
from promptflow._sdk._constants import TRACE_DEFAULT_COLLECTION
from promptflow._sdk._pf_client import PFClient
from promptflow._sdk.entities._trace import Span
from promptflow.tracing import start_trace

TEST_ROOT = (PROMPTFLOW_ROOT / "tests").resolve().absolute()
FLOWS_DIR = (TEST_ROOT / "test_configs/flows").resolve().absolute().as_posix()
FLEX_FLOWS_DIR = (TEST_ROOT / "test_configs/eager_flows").resolve().absolute().as_posix()
PROMPTY_DIR = (TEST_ROOT / "test_configs/prompty").resolve().absolute().as_posix()
DATA_DIR = (TEST_ROOT / "test_configs/datas").resolve().absolute().as_posix()


def load_and_override_span_example(
    trace_id: str,
    span_id: str,
    parent_id: typing.Optional[str],
    line_run_id: str,
) -> typing.Dict:
    # load template span from local example file
    example_span_path = TEST_ROOT / "test_configs/traces/large-data-span-example.json"
    with open(example_span_path, mode="r", encoding="utf-8") as f:
        span_dict = json.load(f)
    # override field(s)
    span_dict["context"]["trace_id"] = trace_id
    span_dict["context"]["span_id"] = span_id
    span_dict["parent_id"] = parent_id
    span_dict["attributes"]["line_run_id"] = line_run_id
    return span_dict


def mock_span(
    trace_id: str,
    span_id: str,
    parent_id: typing.Optional[str],
    line_run_id: str,
) -> Span:
    span_dict = load_and_override_span_example(
        trace_id=trace_id, span_id=span_id, parent_id=parent_id, line_run_id=line_run_id
    )
    # type conversion for timestamp - required for Span constructor
    span_dict["start_time"] = datetime.datetime.fromisoformat(span_dict["start_time"])
    span_dict["end_time"] = datetime.datetime.fromisoformat(span_dict["end_time"])
    # create Span object
    return Span(
        name=span_dict["name"],
        trace_id=trace_id,
        span_id=span_id,
        parent_id=parent_id,
        context=span_dict["context"],
        kind=span_dict["kind"],
        start_time=span_dict["start_time"],
        end_time=span_dict["end_time"],
        status=span_dict["status"],
        attributes=span_dict["attributes"],
        links=span_dict["links"],
        events=span_dict["events"],
        resource=span_dict["resource"],
    )


def mock_span_for_delete_tests(
    run: typing.Optional[str] = None,
    collection: typing.Optional[str] = None,
    start_time: typing.Optional[datetime.datetime] = None,
) -> Span:
    span = mock_span(
        trace_id=str(uuid.uuid4()), span_id=str(uuid.uuid4()), parent_id=None, line_run_id=str(uuid.uuid4())
    )
    if run is not None:
        span.attributes.pop(SpanAttributeFieldName.LINE_RUN_ID)
        span.attributes[SpanAttributeFieldName.BATCH_RUN_ID] = run
        span.attributes[SpanAttributeFieldName.LINE_NUMBER] = 0  # always line 0
    if collection is not None:
        span.resource[SpanResourceFieldName.ATTRIBUTES][SpanResourceAttributesFieldName.COLLECTION] = collection
    if start_time is not None:
        span.start_time = start_time
    span._persist()
    return span


def assert_span_equals(span: Span, expected_span_dict: typing.Dict) -> None:
    span_dict = span._to_rest_object()
    # assert "external_event_data_uris" in span_dict and pop
    assert "external_event_data_uris" in span_dict
    span_dict.pop("external_event_data_uris")
    assert span_dict == expected_span_dict


@pytest.fixture
def collection() -> str:
    _collection = str(uuid.uuid4())
    start_trace(collection=_collection)
    return _collection


@pytest.mark.e2etest
@pytest.mark.sdk_test
class TestTraceEntitiesAndOperations:
    def test_span_to_dict(self) -> None:
        # this should be the groundtruth as OpenTelemetry span spec
        otel_span_path = TEST_ROOT / "test_configs/traces/large-data-span-example.json"
        with open(otel_span_path, mode="r", encoding="utf-8") as f:
            span_dict = json.load(f)
        span_entity = Span(
            name=span_dict["name"],
            trace_id=span_dict["context"]["trace_id"],
            span_id=span_dict["context"]["span_id"],
            parent_id=span_dict["parent_id"],
            context=span_dict["context"],
            kind=span_dict["kind"],
            start_time=datetime.datetime.fromisoformat(span_dict["start_time"]),
            end_time=datetime.datetime.fromisoformat(span_dict["end_time"]),
            status=span_dict["status"],
            attributes=span_dict["attributes"],
            links=span_dict["links"],
            events=span_dict["events"],
            resource=span_dict["resource"],
        )
        otel_span_dict = {
            "name": "openai.resources.chat.completions.Completions.create",
            "context": {
                "trace_id": "32a6fb50e281736543979ce5b929dfdc",
                "span_id": "3a3596a19efef900",
                "trace_state": "",
            },
            "kind": "1",
            "parent_id": "9c63581c6da66596",
            "start_time": "2024-03-21T06:37:22.332582Z",
            "end_time": "2024-03-21T06:37:26.445007Z",
            "status": {
                "status_code": "Ok",
                "description": "",
            },
            "attributes": {
                "framework": "promptflow",
                "span_type": "LLM",
                "function": "openai.resources.chat.completions.Completions.create",
                "node_name": "Azure_OpenAI_GPT_4_Turbo_with_Vision_mrr4",
                "line_run_id": "277fab99-d26e-4c43-8ec4-b0c61669fd68",
                "llm.response.model": "gpt-4",
                "__computed__.cumulative_token_count.completion": "14",
                "__computed__.cumulative_token_count.prompt": "1497",
                "__computed__.cumulative_token_count.total": "1511",
                "llm.usage.completion_tokens": "14",
                "llm.usage.prompt_tokens": "1497",
                "llm.usage.total_tokens": "1511",
            },
            "events": [
                {
                    "name": "promptflow.function.inputs",
                    "timestamp": "2024-03-21T06:37:22.332582Z",
                    "attributes": {
                        "payload": '{"input1": "value1", "input2": "value2"}',
                    },
                },
                {
                    "name": "promptflow.function.output",
                    "timestamp": "2024-03-21T06:37:26.445007Z",
                    "attributes": {
                        "payload": '{"output1": "val1", "output2": "val2"}',
                    },
                },
            ],
            "links": [],
            "resource": {
                "attributes": {
                    "service.name": "promptflow",
                    "collection": "default",
                },
                "schema_url": "",
            },
        }
        assert span_entity.to_dict() == otel_span_dict

    def test_span_persist_and_gets(self, pf: PFClient) -> None:
        trace_id = str(uuid.uuid4())
        span_id = str(uuid.uuid4())
        parent_id = str(uuid.uuid4())
        line_run_id = str(uuid.uuid4())
        span = mock_span(trace_id=trace_id, span_id=span_id, parent_id=parent_id, line_run_id=line_run_id)
        span._persist()
        # trace operations - get span
        # eager load
        eager_load_span = pf.traces.get_span(trace_id=trace_id, span_id=span_id, lazy_load=False)
        expected_span_dict = load_and_override_span_example(
            trace_id=trace_id, span_id=span_id, parent_id=parent_id, line_run_id=line_run_id
        )
        assert_span_equals(eager_load_span, expected_span_dict)
        # lazy load (default)
        lazy_load_span = pf.traces.get_span(trace_id=trace_id, span_id=span_id)
        # events.attributes should be empty in lazy load mode
        for i in range(len(expected_span_dict["events"])):
            expected_span_dict["events"][i]["attributes"] = dict()
        assert_span_equals(lazy_load_span, expected_span_dict)

    def test_aggregation_node_in_eval_run(self, pf: PFClient) -> None:
        # mock a span generated from an aggregation node in an eval run
        # whose attributes has `referenced.batch_run_id`, no `line_number`
        span = mock_span(
            trace_id=str(uuid.uuid4()),
            span_id=str(uuid.uuid4()),
            parent_id=None,
            line_run_id=str(uuid.uuid4()),
        )
        batch_run_id = str(uuid.uuid4())
        span.attributes.pop(SpanAttributeFieldName.LINE_RUN_ID)
        span.attributes[SpanAttributeFieldName.BATCH_RUN_ID] = batch_run_id
        span.attributes[SpanAttributeFieldName.REFERENCED_BATCH_RUN_ID] = str(uuid.uuid4())
        span._persist()
        # list and assert to ensure the persist is successful
        line_runs = pf.traces.list_line_runs(runs=[batch_run_id])
        assert len(line_runs) == 1

    def test_spans_persist_and_line_run_gets(self, pf: PFClient) -> None:
        trace_id = str(uuid.uuid4())
        non_root_span_id = str(uuid.uuid4())
        root_span_id = str(uuid.uuid4())
        line_run_id = str(uuid.uuid4())
        # non-root span
        span = mock_span(
            trace_id=trace_id,
            span_id=non_root_span_id,
            parent_id=root_span_id,
            line_run_id=line_run_id,
        )
        span._persist()
        running_line_run = pf.traces.get_line_run(line_run_id=line_run_id)
        expected_running_line_run_dict = {
            "line_run_id": line_run_id,
            "trace_id": trace_id,
            "root_span_id": None,
            "inputs": None,
            "outputs": None,
            "start_time": "2024-03-21T06:37:22.332582",
            "end_time": None,
            "status": RUNNING_LINE_RUN_STATUS,
            "duration": None,
            "name": None,
            "kind": None,
            "collection": TRACE_DEFAULT_COLLECTION,
            "cumulative_token_count": None,
            "parent_id": None,
            "run": None,
            "line_number": None,
            "experiment": None,
            "session_id": None,
            "evaluations": None,
        }
        assert running_line_run._to_rest_object() == expected_running_line_run_dict
        # root span
        span = mock_span(
            trace_id=trace_id,
            span_id=root_span_id,
            parent_id=None,
            line_run_id=line_run_id,
        )
        span._persist()
        terminated_line_run = pf.traces.get_line_run(line_run_id=line_run_id)
        expected_terminated_line_run_dict = {
            "line_run_id": line_run_id,
            "trace_id": trace_id,
            "root_span_id": root_span_id,
            "inputs": {"input1": "value1", "input2": "value2"},
            "outputs": {"output1": "val1", "output2": "val2"},
            "start_time": "2024-03-21T06:37:22.332582",
            "end_time": "2024-03-21T06:37:26.445007",
            "status": "Ok",
            "duration": 4.112425,
            "name": "openai.resources.chat.completions.Completions.create",
            "kind": "LLM",
            "collection": TRACE_DEFAULT_COLLECTION,
            "cumulative_token_count": {
                "completion": 14,
                "prompt": 1497,
                "total": 1511,
            },
            "parent_id": None,
            "run": None,
            "line_number": None,
            "experiment": None,
            "session_id": None,
            "evaluations": None,
        }
        assert terminated_line_run._to_rest_object() == expected_terminated_line_run_dict

    def test_span_io_in_attrs_persist(self, pf: PFClient) -> None:
        trace_id, span_id, line_run_id = str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4())
        span = mock_span(trace_id=trace_id, span_id=span_id, parent_id=None, line_run_id=line_run_id)
        # empty span.events and move inputs/output to span.attributes
        inputs = {"input1": "value1", "input2": "value2"}
        output = {"output1": "val1", "output2": "val2"}
        span.attributes[SpanAttributeFieldName.INPUTS] = json.dumps(inputs)
        span.attributes[SpanAttributeFieldName.OUTPUT] = json.dumps(output)
        span.events = list()
        span._persist()
        line_run = pf.traces.get_line_run(line_run_id=line_run_id)
        assert line_run.inputs == inputs
        assert line_run.outputs == output

    def test_span_non_json_io_in_attrs_persist(self, pf: PFClient) -> None:
        trace_id, span_id, line_run_id = str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4())
        span = mock_span(trace_id=trace_id, span_id=span_id, parent_id=None, line_run_id=line_run_id)
        # empty span.events and set non-JSON inputs/output to span.attributes
        inputs = {"input1": "value1", "input2": "value2"}
        output = {"output1": "val1", "output2": "val2"}
        span.attributes[SpanAttributeFieldName.INPUTS] = str(inputs)
        span.attributes[SpanAttributeFieldName.OUTPUT] = str(output)
        span.events = list()
        span._persist()
        line_run = pf.traces.get_line_run(line_run_id=line_run_id)
        assert isinstance(line_run.inputs, str) and line_run.inputs == str(inputs)
        assert isinstance(line_run.outputs, str) and line_run.outputs == str(output)

    def test_span_with_nan_as_io(self, pf: PFClient) -> None:
        trace_id, span_id, line_run_id = str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4())
        span = mock_span(trace_id=trace_id, span_id=span_id, parent_id=None, line_run_id=line_run_id)
        span.events[0]["attributes"]["payload"] = json.dumps(dict(input1=float("nan"), input2=float("inf")))
        span.events[1]["attributes"]["payload"] = json.dumps(dict(output1=float("nan"), output2=float("-inf")))
        span._persist()
        line_run = pf.traces.get_line_run(line_run_id=line_run_id)
        line_run_inputs, line_run_outputs = line_run.inputs, line_run.outputs
        assert isinstance(line_run_inputs["input1"], str) and line_run_inputs["input1"] == "NaN"
        assert isinstance(line_run_inputs["input2"], str) and line_run_inputs["input2"] == "Infinity"
        assert isinstance(line_run_outputs["output1"], str) and line_run_outputs["output1"] == "NaN"
        assert isinstance(line_run_outputs["output2"], str) and line_run_outputs["output2"] == "-Infinity"

    def test_delete_traces_three_tables(self, pf: PFClient) -> None:
        # trace operation does not expose API for events and spans
        # so directly use ORM class to list and assert events and spans existence and deletion
        from promptflow._sdk._orm.trace import Event as ORMEvent
        from promptflow._sdk._orm.trace import LineRun as ORMLineRun
        from promptflow._sdk._orm.trace import Span as ORMSpan

        mock_run = str(uuid.uuid4())
        mock_span = mock_span_for_delete_tests(run=mock_run)
        # assert events, span and line_run are persisted
        assert len(ORMEvent.list(trace_id=mock_span.trace_id, span_id=mock_span.span_id)) == 2
        assert len(ORMSpan.list(trace_ids=[mock_span.trace_id])) == 1
        assert len(ORMLineRun.list(runs=[mock_run])) == 1
        # delete traces and assert all traces are deleted
        pf.traces.delete(run=mock_run)
        assert len(ORMEvent.list(trace_id=mock_span.trace_id, span_id=mock_span.span_id)) == 0
        assert len(ORMSpan.list(trace_ids=[mock_span.trace_id])) == 0
        assert len(ORMLineRun.list(runs=[mock_run])) == 0

    def test_delete_traces_with_run(self, pf: PFClient) -> None:
        mock_run = str(uuid.uuid4())
        mock_span_for_delete_tests(run=mock_run)
        assert len(pf.traces.list_line_runs(runs=[mock_run])) == 1
        pf.traces.delete(run=mock_run)
        assert len(pf.traces.list_line_runs(runs=[mock_run])) == 0

    def test_delete_traces_with_collection(self, pf: PFClient) -> None:
        mock_collection = str(uuid.uuid4())
        mock_span_for_delete_tests(collection=mock_collection)
        assert len(pf.traces.list_line_runs(collection=mock_collection)) == 1
        pf.traces.delete(collection=mock_collection)
        assert len(pf.traces.list_line_runs(collection=mock_collection)) == 0

    def test_delete_traces_with_collection_and_started_before(self, pf: PFClient) -> None:
        # mock some traces that start 2 days before, and delete those start 1 days before
        mock_start_time = datetime.datetime.now() - datetime.timedelta(days=2)
        collection1, collection2 = str(uuid.uuid4()), str(uuid.uuid4())
        mock_span_for_delete_tests(collection=collection1, start_time=mock_start_time)
        mock_span_for_delete_tests(collection=collection2, start_time=mock_start_time)
        assert (
            len(pf.traces.list_line_runs(collection=collection1)) == 1
            and len(pf.traces.list_line_runs(collection=collection2)) == 1
        )
        delete_query_time = datetime.datetime.now() - datetime.timedelta(days=1)
        pf.traces.delete(collection=collection1, started_before=delete_query_time.isoformat())
        # only collection1 traces are deleted
        assert (
            len(pf.traces.list_line_runs(collection=collection1)) == 0
            and len(pf.traces.list_line_runs(collection=collection2)) == 1
        )
        pf.traces.delete(collection=collection2, started_before=delete_query_time.isoformat())
        assert len(pf.traces.list_line_runs(collection=collection2)) == 0

    def test_delete_traces_dry_run(self, pf: PFClient) -> None:
        mock_run = str(uuid.uuid4())
        mock_span_for_delete_tests(run=mock_run)
        num_traces = pf.traces.delete(run=mock_run, dry_run=True)
        assert num_traces == 1

    def test_basic_search_line_runs(self, pf: PFClient) -> None:
        trace_id = str(uuid.uuid4())
        span_id = str(uuid.uuid4())
        line_run_id = str(uuid.uuid4())
        span = mock_span(trace_id=trace_id, span_id=span_id, parent_id=None, line_run_id=line_run_id)
        name = str(uuid.uuid4())
        span.name = name
        span._persist()
        expr = f"name == '{name}'"
        line_runs = pf.traces._search_line_runs(expression=expr)
        assert len(line_runs) == 1

    @pytest.mark.skipif(
        platform.system() == "Windows" and sys.version_info < (3, 9),
        reason="Python 3.9+ is required on Windows to support json_extract",
    )
    def test_search_line_runs_with_tokens(self, pf: PFClient) -> None:
        num_line_runs = 5
        trace_ids = list()
        name = str(uuid.uuid4())
        for _ in range(num_line_runs):
            trace_id = str(uuid.uuid4())
            span_id = str(uuid.uuid4())
            line_run_id = str(uuid.uuid4())
            span = mock_span(trace_id=trace_id, span_id=span_id, parent_id=None, line_run_id=line_run_id)
            span.name = name
            span.attributes.update({"__computed__.cumulative_token_count.total": "42"})
            span._persist()
            trace_ids.append(trace_id)
        expr = f"name == '{name}' and total < 100"
        line_runs = pf.traces._search_line_runs(expression=expr)
        assert len(line_runs) == num_line_runs
        # assert these line runs are exactly the ones we just persisted
        line_run_trace_ids = {line_run.trace_id for line_run in line_runs}
        assert len(set(trace_ids) & line_run_trace_ids) == num_line_runs

    def test_list_collection(self, pf: PFClient) -> None:
        collection = str(uuid.uuid4())
        span = mock_span(
            trace_id=str(uuid.uuid4()), span_id=str(uuid.uuid4()), parent_id=None, line_run_id=str(uuid.uuid4())
        )
        # make span start time a week later, so that it can be the latest collection
        span.start_time = datetime.datetime.now() + datetime.timedelta(days=7)
        span.start_time = datetime.datetime.now() + datetime.timedelta(days=8)
        span.resource[SpanResourceFieldName.ATTRIBUTES][SpanResourceAttributesFieldName.COLLECTION] = collection
        span._persist()
        collections = pf.traces._list_collections(limit=1)
        assert len(collections) == 1 and collections[0].name == collection

    def test_list_collection_with_time_priority(self, pf: PFClient) -> None:
        collection1, collection2 = str(uuid.uuid4()), str(uuid.uuid4())
        for collection in (collection1, collection2):
            span = mock_span(
                trace_id=str(uuid.uuid4()), span_id=str(uuid.uuid4()), parent_id=None, line_run_id=str(uuid.uuid4())
            )
            # make span start time a week later, so that it can be the latest collection
            span.start_time = datetime.datetime.now() + datetime.timedelta(days=7)
            span.start_time = datetime.datetime.now() + datetime.timedelta(days=8)
            span.resource[SpanResourceFieldName.ATTRIBUTES][SpanResourceAttributesFieldName.COLLECTION] = collection
            span._persist()
            # sleep 1 second to ensure the second span is later than the first
            time.sleep(1)
        collections = pf.traces._list_collections(limit=1)
        assert len(collections) == 1 and collections[0].name == collection2
        collections = pf.traces._list_collections(limit=2)
        assert len(collections) == 2 and collections[1].name == collection1


@pytest.mark.usefixtures("use_secrets_config_file", "recording_injection", "setup_local_connection")
@pytest.mark.e2etest
@pytest.mark.sdk_test
class TestTraceWithDevKit:
    def test_flow_test_trace_enabled(self, pf: PFClient) -> None:
        import promptflow._sdk._orchestrator.test_submitter

        with mock.patch("promptflow._sdk._configuration.Configuration.is_internal_features_enabled") as mock_func:
            mock_func.return_value = True
            with patch.object(promptflow._sdk._orchestrator.test_submitter, "start_trace") as mock_start_trace:
                inputs = {"url": "https://www.youtube.com/watch?v=o5ZQyXaAv1g", "answer": "Channel", "evidence": "Url"}
                pf.test(flow=Path(f"{FLOWS_DIR}/web_classification").absolute(), inputs=inputs)
                assert mock_start_trace.call_count == 1

    def test_flow_test_single_node_trace_not_enabled(self, pf: PFClient) -> None:
        import promptflow._sdk._orchestrator.test_submitter

        with mock.patch("promptflow._sdk._configuration.Configuration.is_internal_features_enabled") as mock_func:
            mock_func.return_value = True
            with patch.object(promptflow._sdk._orchestrator.test_submitter, "start_trace") as mock_start_trace:
                pf.test(
                    flow=Path(f"{FLOWS_DIR}/web_classification").absolute(),
                    inputs={"fetch_url": "https://www.youtube.com/watch?v=o5ZQyXaAv1g"},
                    node="fetch_text_content_from_url",
                )
                assert mock_start_trace.call_count == 0


@pytest.mark.usefixtures("otlp_collector", "recording_injection", "setup_local_connection", "use_secrets_config_file")
@pytest.mark.e2etest
@pytest.mark.sdk_test
class TestTraceLifeCycle:
    """End-to-end tests that cover the trace lifecycle."""

    def _clear_module_cache(self, module_name) -> None:
        # referenced from test_flow_test.py::clear_module_cache
        try:
            del sys.modules[module_name]
        except Exception:  # pylint: disable=broad-except
            pass

    def _pf_test_and_assert(
        self,
        pf: PFClient,
        flow_path: Path,
        inputs: typing.Dict[str, str],
        collection: str,
    ) -> None:
        pf.test(flow=flow_path, inputs=inputs)
        line_runs = pf.traces.list_line_runs(collection=collection)
        assert len(line_runs) == 1

    def test_flow_test_dag_flow(self, pf: PFClient, collection: str) -> None:
        flow_path = Path(f"{FLOWS_DIR}/web_classification").absolute()
        inputs = {"url": "https://www.youtube.com/watch?v=o5ZQyXaAv1g", "answer": "Channel", "evidence": "Url"}
        self._pf_test_and_assert(pf, flow_path, inputs, collection)

    def test_flow_test_flex_flow(self, pf: PFClient, collection: str) -> None:
        self._clear_module_cache("entry")
        flow_path = Path(f"{FLEX_FLOWS_DIR}/simple_with_yaml").absolute()
        inputs = {"input_val": "val1"}
        self._pf_test_and_assert(pf, flow_path, inputs, collection)

    def test_flow_test_prompty(self, pf: PFClient, collection: str) -> None:
        flow_path = Path(f"{PROMPTY_DIR}/prompty_example.prompty").absolute()
        inputs = {"question": "what is the result of 1+1?"}
        self._pf_test_and_assert(pf, flow_path, inputs, collection)

    def _pf_run_and_assert(
        self,
        pf: PFClient,
        flow_path: Path,
        data_path: Path,
        expected_number_lines: int,
    ):
        run = pf.run(flow=flow_path, data=data_path)
        line_runs = pf.traces.list_line_runs(runs=run.name)
        assert len(line_runs) == expected_number_lines

    def test_batch_run_dag_flow(self, pf: PFClient) -> None:
        flow_path = Path(f"{FLOWS_DIR}/web_classification").absolute()
        data_path = Path(f"{DATA_DIR}/webClassification3.jsonl").absolute()
        self._pf_run_and_assert(pf, flow_path, data_path, expected_number_lines=3)

    def test_batch_run_flex_flow(self, pf: PFClient) -> None:
        flow_path = Path(f"{FLEX_FLOWS_DIR}/simple_with_yaml").absolute()
        data_path = Path(f"{DATA_DIR}/simple_eager_flow_data.jsonl").absolute()
        self._pf_run_and_assert(pf, flow_path, data_path, expected_number_lines=1)

    def test_batch_run_prompty(self, pf: PFClient) -> None:
        flow_path = Path(f"{PROMPTY_DIR}/prompty_example.prompty").absolute()
        data_path = Path(f"{DATA_DIR}/prompty_inputs.jsonl").absolute()
        self._pf_run_and_assert(pf, flow_path, data_path, expected_number_lines=3)