From 5ca3da92a599b49c62a4d242da5ef29a808e0415 Mon Sep 17 00:00:00 2001 From: Maksim Moiseenkov Date: Mon, 11 Sep 2023 14:17:24 +0000 Subject: [PATCH] Add explicit support of stream (realtime) pipelines for CloudDataFusionStartPipelineOperator --- .../google/cloud/hooks/datafusion.py | 48 ++++++++-- .../google/cloud/operators/datafusion.py | 7 ++ .../google/cloud/triggers/datafusion.py | 6 ++ .../google/cloud/utils/datafusion.py | 33 +++++++ .../google/cloud/hooks/test_datafusion.py | 90 ++++++++++++++++++- .../google/cloud/operators/test_datafusion.py | 5 ++ .../google/cloud/triggers/test_datafusion.py | 3 + .../google/cloud/utils/test_datafusion.py | 37 ++++++++ 8 files changed, 219 insertions(+), 10 deletions(-) create mode 100644 airflow/providers/google/cloud/utils/datafusion.py create mode 100644 tests/providers/google/cloud/utils/test_datafusion.py diff --git a/airflow/providers/google/cloud/hooks/datafusion.py b/airflow/providers/google/cloud/hooks/datafusion.py index 54ae768f15195..b3cfd6ee9b4fc 100644 --- a/airflow/providers/google/cloud/hooks/datafusion.py +++ b/airflow/providers/google/cloud/hooks/datafusion.py @@ -31,6 +31,7 @@ from googleapiclient.discovery import Resource, build from airflow.exceptions import AirflowException, AirflowNotFoundException +from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType from airflow.providers.google.common.hooks.base_google import ( PROVIDE_PROJECT_ID, GoogleBaseAsyncHook, @@ -105,6 +106,7 @@ def wait_for_pipeline_state( pipeline_name: str, pipeline_id: str, instance_url: str, + pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH, namespace: str = "default", success_states: list[str] | None = None, failure_states: list[str] | None = None, @@ -120,6 +122,7 @@ def wait_for_pipeline_state( workflow = self.get_pipeline_workflow( pipeline_name=pipeline_name, pipeline_id=pipeline_id, + pipeline_type=pipeline_type, instance_url=instance_url, namespace=namespace, ) @@ -432,13 +435,14 @@ def get_pipeline_workflow( pipeline_name: str, instance_url: str, pipeline_id: str, + pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH, namespace: str = "default", ) -> Any: url = os.path.join( self._base_url(instance_url, namespace), quote(pipeline_name), - "workflows", - "DataPipelineWorkflow", + f"{self.cdap_program_type(pipeline_type=pipeline_type)}s", + self.cdap_program_id(pipeline_type=pipeline_type), "runs", quote(pipeline_id), ) @@ -453,6 +457,7 @@ def start_pipeline( self, pipeline_name: str, instance_url: str, + pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH, namespace: str = "default", runtime_args: dict[str, Any] | None = None, ) -> str: @@ -460,6 +465,7 @@ def start_pipeline( Starts a Cloud Data Fusion pipeline. Works for both batch and stream pipelines. :param pipeline_name: Your pipeline name. + :param pipeline_type: Optional pipeline type (BATCH by default). :param instance_url: Endpoint on which the REST APIs is accessible for the instance. :param runtime_args: Optional runtime JSON args to be passed to the pipeline :param namespace: if your pipeline belongs to a Basic edition instance, the namespace ID @@ -480,9 +486,9 @@ def start_pipeline( body = [ { "appId": pipeline_name, - "programType": "workflow", - "programId": "DataPipelineWorkflow", "runtimeargs": runtime_args, + "programType": self.cdap_program_type(pipeline_type=pipeline_type), + "programId": self.cdap_program_id(pipeline_type=pipeline_type), } ] response = self._cdap_request(url=url, method="POST", body=body) @@ -514,6 +520,30 @@ def stop_pipeline(self, pipeline_name: str, instance_url: str, namespace: str = response, f"Stopping a pipeline failed with code {response.status}" ) + @staticmethod + def cdap_program_type(pipeline_type: DataFusionPipelineType) -> str: + """Retrieves CDAP Program type depending on the pipeline type. + + :param pipeline_type: Pipeline type. + """ + program_types = { + DataFusionPipelineType.BATCH: "workflow", + DataFusionPipelineType.STREAM: "spark", + } + return program_types.get(pipeline_type, "") + + @staticmethod + def cdap_program_id(pipeline_type: DataFusionPipelineType) -> str: + """Retrieves CDAP Program id depending on the pipeline type. + + :param pipeline_type: Pipeline type. + """ + program_ids = { + DataFusionPipelineType.BATCH: "DataPipelineWorkflow", + DataFusionPipelineType.STREAM: "DataStreamsSparkStreaming", + } + return program_ids.get(pipeline_type, "") + class DataFusionAsyncHook(GoogleBaseAsyncHook): """Class to get asynchronous hook for DataFusion.""" @@ -561,10 +591,13 @@ async def get_pipeline( pipeline_name: str, pipeline_id: str, session, + pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH, ): + program_type = self.sync_hook_class.cdap_program_type(pipeline_type=pipeline_type) + program_id = self.sync_hook_class.cdap_program_id(pipeline_type=pipeline_type) base_url_link = self._base_url(instance_url, namespace) url = urljoin( - base_url_link, f"{quote(pipeline_name)}/workflows/DataPipelineWorkflow/runs/{quote(pipeline_id)}" + base_url_link, f"{quote(pipeline_name)}/{program_type}s/{program_id}/runs/{quote(pipeline_id)}" ) return await self._get_link(url=url, session=session) @@ -573,6 +606,7 @@ async def get_pipeline_status( pipeline_name: str, instance_url: str, pipeline_id: str, + pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH, namespace: str = "default", success_states: list[str] | None = None, ) -> str: @@ -581,7 +615,8 @@ async def get_pipeline_status( :param pipeline_name: Your pipeline name. :param instance_url: Endpoint on which the REST APIs is accessible for the instance. - :param pipeline_id: Unique pipeline ID associated with specific pipeline + :param pipeline_id: Unique pipeline ID associated with specific pipeline. + :param pipeline_type: Optional pipeline type (by default batch). :param namespace: if your pipeline belongs to a Basic edition instance, the namespace ID is always default. If your pipeline belongs to an Enterprise edition instance, you can create a namespace. @@ -596,6 +631,7 @@ async def get_pipeline_status( namespace=namespace, pipeline_name=pipeline_name, pipeline_id=pipeline_id, + pipeline_type=pipeline_type, session=session, ) pipeline = await pipeline.json(content_type=None) diff --git a/airflow/providers/google/cloud/operators/datafusion.py b/airflow/providers/google/cloud/operators/datafusion.py index d82a27087114e..d63de8442dc70 100644 --- a/airflow/providers/google/cloud/operators/datafusion.py +++ b/airflow/providers/google/cloud/operators/datafusion.py @@ -33,6 +33,7 @@ ) from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator from airflow.providers.google.cloud.triggers.datafusion import DataFusionStartPipelineTrigger +from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType if TYPE_CHECKING: from airflow.utils.context import Context @@ -708,6 +709,7 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator): :ref:`howto/operator:CloudDataFusionStartPipelineOperator` :param pipeline_name: Your pipeline name. + :param pipeline_type: Optional pipeline type (BATCH by default). :param instance_name: The name of the instance. :param success_states: If provided the operator will wait for pipeline to be in one of the provided states. @@ -752,6 +754,7 @@ def __init__( pipeline_name: str, instance_name: str, location: str, + pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH, runtime_args: dict[str, Any] | None = None, success_states: list[str] | None = None, namespace: str = "default", @@ -767,6 +770,7 @@ def __init__( ) -> None: super().__init__(**kwargs) self.pipeline_name = pipeline_name + self.pipeline_type = pipeline_type self.runtime_args = runtime_args self.namespace = namespace self.instance_name = instance_name @@ -800,6 +804,7 @@ def execute(self, context: Context) -> str: api_url = instance["apiEndpoint"] pipeline_id = hook.start_pipeline( pipeline_name=self.pipeline_name, + pipeline_type=self.pipeline_type, instance_url=api_url, namespace=self.namespace, runtime_args=self.runtime_args, @@ -824,6 +829,7 @@ def execute(self, context: Context) -> str: instance_url=api_url, namespace=self.namespace, pipeline_name=self.pipeline_name, + pipeline_type=self.pipeline_type.value, pipeline_id=pipeline_id, poll_interval=self.poll_interval, gcp_conn_id=self.gcp_conn_id, @@ -839,6 +845,7 @@ def execute(self, context: Context) -> str: success_states=self.success_states, pipeline_id=pipeline_id, pipeline_name=self.pipeline_name, + pipeline_type=self.pipeline_type, namespace=self.namespace, instance_url=api_url, timeout=self.pipeline_timeout, diff --git a/airflow/providers/google/cloud/triggers/datafusion.py b/airflow/providers/google/cloud/triggers/datafusion.py index 06bf5e053eab2..d1419bfbc5f5d 100644 --- a/airflow/providers/google/cloud/triggers/datafusion.py +++ b/airflow/providers/google/cloud/triggers/datafusion.py @@ -20,6 +20,7 @@ from typing import Any, AsyncIterator, Sequence from airflow.providers.google.cloud.hooks.datafusion import DataFusionAsyncHook +from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType from airflow.triggers.base import BaseTrigger, TriggerEvent @@ -30,6 +31,7 @@ class DataFusionStartPipelineTrigger(BaseTrigger): :param pipeline_name: Your pipeline name. :param instance_url: Endpoint on which the REST APIs is accessible for the instance. :param pipeline_id: Unique pipeline ID associated with specific pipeline + :param pipeline_type: Your pipeline type. :param namespace: if your pipeline belongs to a Basic edition instance, the namespace ID is always default. If your pipeline belongs to an Enterprise edition instance, you can create a namespace. @@ -51,6 +53,7 @@ def __init__( namespace: str, pipeline_name: str, pipeline_id: str, + pipeline_type: str, poll_interval: float = 3.0, gcp_conn_id: str = "google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, @@ -61,6 +64,7 @@ def __init__( self.namespace = namespace self.pipeline_name = pipeline_name self.pipeline_id = pipeline_id + self.pipeline_type = pipeline_type self.poll_interval = poll_interval self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain @@ -76,6 +80,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]: "namespace": self.namespace, "pipeline_name": self.pipeline_name, "pipeline_id": self.pipeline_id, + "pipeline_type": self.pipeline_type, "success_states": self.success_states, }, ) @@ -92,6 +97,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override] namespace=self.namespace, pipeline_name=self.pipeline_name, pipeline_id=self.pipeline_id, + pipeline_type=DataFusionPipelineType.from_str(self.pipeline_type), ) if response_from_hook == "success": yield TriggerEvent( diff --git a/airflow/providers/google/cloud/utils/datafusion.py b/airflow/providers/google/cloud/utils/datafusion.py new file mode 100644 index 0000000000000..20101392a9204 --- /dev/null +++ b/airflow/providers/google/cloud/utils/datafusion.py @@ -0,0 +1,33 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from enum import Enum + + +class DataFusionPipelineType(Enum): + """Enum for Data Fusion pipeline types.""" + + BATCH = "batch" + STREAM = "stream" + + @staticmethod + def from_str(value: str) -> DataFusionPipelineType: + value_to_item = {item.value: item for item in DataFusionPipelineType} + if value in value_to_item: + return value_to_item[value] + raise ValueError(f"Invalid value '{value}'. Valid values are: {[i for i in value_to_item.keys()]}") diff --git a/tests/providers/google/cloud/hooks/test_datafusion.py b/tests/providers/google/cloud/hooks/test_datafusion.py index 9ee4b9d39e43c..d6815da5b8e7b 100644 --- a/tests/providers/google/cloud/hooks/test_datafusion.py +++ b/tests/providers/google/cloud/hooks/test_datafusion.py @@ -26,6 +26,7 @@ from airflow import AirflowException from airflow.providers.google.cloud.hooks.datafusion import DataFusionAsyncHook, DataFusionHook +from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType from tests.providers.google.cloud.utils.base_gcp_mock import mock_base_gcp_hook_default_project_id API_VERSION = "v1beta1" @@ -50,6 +51,11 @@ f"googleusercontent.com/api/v3/namespaces/{NAMESPACE}/apps/{PIPELINE_NAME}" f"/workflows/DataPipelineWorkflow/runs/{PIPELINE_ID}" ) +CONSTRUCTED_PIPELINE_STREAM_URL_GET = ( + f"https://{INSTANCE_NAME}-{PROJECT_ID}-dot-eun1.datafusion." + f"googleusercontent.com/api/v3/namespaces/{NAMESPACE}/apps/{PIPELINE_NAME}" + f"/apsrkss/DataStreamsSparkStreaming/runs/{PIPELINE_ID}" +) class MockResponse: @@ -350,6 +356,29 @@ def test_start_pipeline(self, mock_request, hook): url=f"{INSTANCE_URL}/v3/namespaces/default/start", method="POST", body=body ) + @mock.patch(HOOK_STR.format("DataFusionHook._cdap_request")) + def test_start_pipeline_stream(self, mock_request, hook): + run_id = 1234 + mock_request.return_value = mock.MagicMock(status=200, data=f'[{{"runId":{run_id}}}]') + + hook.start_pipeline( + pipeline_name=PIPELINE_NAME, + instance_url=INSTANCE_URL, + runtime_args=RUNTIME_ARGS, + pipeline_type=DataFusionPipelineType.STREAM, + ) + body = [ + { + "appId": PIPELINE_NAME, + "programType": "spark", + "programId": "DataStreamsSparkStreaming", + "runtimeargs": RUNTIME_ARGS, + } + ] + mock_request.assert_called_once_with( + url=f"{INSTANCE_URL}/v3/namespaces/default/start", method="POST", body=body + ) + @mock.patch(HOOK_STR.format("DataFusionHook._cdap_request")) def test_start_pipeline_should_fail_if_empty_data_response(self, mock_request, hook): mock_request.return_value.status = 200 @@ -443,6 +472,22 @@ def test_get_pipeline_workflow(self, mock_request, hook): method="GET", ) + @mock.patch(HOOK_STR.format("DataFusionHook._cdap_request")) + def test_get_pipeline_workflow_stream(self, mock_request, hook): + run_id = 1234 + mock_request.return_value = mock.MagicMock(status=200, data=f'[{{"runId":{run_id}}}]') + hook.get_pipeline_workflow( + pipeline_name=PIPELINE_NAME, + instance_url=INSTANCE_URL, + pipeline_id=PIPELINE_ID, + pipeline_type=DataFusionPipelineType.STREAM, + ) + mock_request.assert_called_once_with( + url=f"{INSTANCE_URL}/v3/namespaces/default/apps/{PIPELINE_NAME}/" + f"sparks/DataStreamsSparkStreaming/runs/{PIPELINE_ID}", + method="GET", + ) + @mock.patch(HOOK_STR.format("DataFusionHook._cdap_request")) def test_get_pipeline_workflow_should_fail_if_empty_data_response(self, mock_request, hook): mock_request.return_value.status = 200 @@ -474,6 +519,28 @@ def test_get_pipeline_workflow_should_fail_if_status_not_200(self, mock_request, method="GET", ) + @pytest.mark.parametrize( + "pipeline_type, expected_program_type", + [ + (DataFusionPipelineType.BATCH, "workflow"), + (DataFusionPipelineType.STREAM, "spark"), + ("non existing value", ""), + ], + ) + def test_cdap_program_type(self, pipeline_type, expected_program_type): + assert DataFusionHook.cdap_program_type(pipeline_type) == expected_program_type + + @pytest.mark.parametrize( + "pipeline_type, expected_program_id", + [ + (DataFusionPipelineType.BATCH, "DataPipelineWorkflow"), + (DataFusionPipelineType.STREAM, "DataStreamsSparkStreaming"), + ("non existing value", ""), + ], + ) + def test_cdap_program_id(self, pipeline_type, expected_program_id): + assert DataFusionHook.cdap_program_id(pipeline_type) == expected_program_id + class TestDataFusionHookAsynch: def test_delegate_to_runtime_error(self): @@ -493,13 +560,20 @@ async def test_async_get_pipeline_should_execute_successfully(self, mocked_link, mocked_link.assert_awaited_once_with(url=CONSTRUCTED_PIPELINE_URL, session=session) @pytest.mark.asyncio + @pytest.mark.parametrize( + "pipeline_type, constructed_url", + [ + (DataFusionPipelineType.BATCH, CONSTRUCTED_PIPELINE_URL_GET), + (DataFusionPipelineType.STREAM, CONSTRUCTED_PIPELINE_STREAM_URL_GET), + ], + ) @mock.patch(HOOK_STR.format("DataFusionAsyncHook.get_pipeline")) async def test_async_get_pipeline_status_completed_should_execute_successfully( - self, mocked_get, hook_async + self, mocked_get, hook_async, pipeline_type, constructed_url ): response = aiohttp.ClientResponse( "get", - URL(CONSTRUCTED_PIPELINE_URL_GET), + URL(constructed_url), request_info=mock.Mock(), writer=mock.Mock(), continue100=None, @@ -523,14 +597,21 @@ async def test_async_get_pipeline_status_completed_should_execute_successfully( assert pipeline_status == "success" @pytest.mark.asyncio + @pytest.mark.parametrize( + "pipeline_type, constructed_url", + [ + (DataFusionPipelineType.BATCH, CONSTRUCTED_PIPELINE_URL_GET), + (DataFusionPipelineType.STREAM, CONSTRUCTED_PIPELINE_STREAM_URL_GET), + ], + ) @mock.patch(HOOK_STR.format("DataFusionAsyncHook.get_pipeline")) async def test_async_get_pipeline_status_running_should_execute_successfully( - self, mocked_get, hook_async + self, mocked_get, hook_async, pipeline_type, constructed_url ): """Assets that the DataFusionAsyncHook returns pending response when job is still in running state""" response = aiohttp.ClientResponse( "get", - URL(CONSTRUCTED_PIPELINE_URL_GET), + URL(constructed_url), request_info=mock.Mock(), writer=mock.Mock(), continue100=None, @@ -548,6 +629,7 @@ async def test_async_get_pipeline_status_running_should_execute_successfully( pipeline_name=PIPELINE_NAME, instance_url=INSTANCE_URL, pipeline_id=PIPELINE_ID, + pipeline_type=pipeline_type, namespace=NAMESPACE, ) mocked_get.assert_awaited_once() diff --git a/tests/providers/google/cloud/operators/test_datafusion.py b/tests/providers/google/cloud/operators/test_datafusion.py index fd5ed2d648dd4..988a64e9ab21c 100644 --- a/tests/providers/google/cloud/operators/test_datafusion.py +++ b/tests/providers/google/cloud/operators/test_datafusion.py @@ -36,6 +36,7 @@ CloudDataFusionUpdateInstanceOperator, ) from airflow.providers.google.cloud.triggers.datafusion import DataFusionStartPipelineTrigger +from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType HOOK_STR = "airflow.providers.google.cloud.operators.datafusion.DataFusionHook" @@ -235,12 +236,14 @@ def test_execute_check_hook_call_should_execute_successfully(self, mock_hook): pipeline_name=PIPELINE_NAME, namespace=NAMESPACE, runtime_args=RUNTIME_ARGS, + pipeline_type=DataFusionPipelineType.BATCH, ) mock_hook.return_value.wait_for_pipeline_state.assert_called_once_with( success_states=[*SUCCESS_STATES, PipelineStates.RUNNING], pipeline_id=PIPELINE_ID, pipeline_name=PIPELINE_NAME, + pipeline_type=DataFusionPipelineType.BATCH, namespace=NAMESPACE, instance_url=INSTANCE_URL, timeout=300, @@ -275,6 +278,7 @@ def test_execute_check_hook_call_asynch_param_should_execute_successfully(self, pipeline_name=PIPELINE_NAME, namespace=NAMESPACE, runtime_args=RUNTIME_ARGS, + pipeline_type=DataFusionPipelineType.BATCH, ) mock_hook.return_value.wait_for_pipeline_state.assert_not_called() @@ -374,6 +378,7 @@ def test_asynch_execute_check_hook_call_should_execute_successfully(self, mock_h pipeline_name=PIPELINE_NAME, namespace=NAMESPACE, runtime_args=RUNTIME_ARGS, + pipeline_type=DataFusionPipelineType.BATCH, ) @mock.patch(HOOK_STR) diff --git a/tests/providers/google/cloud/triggers/test_datafusion.py b/tests/providers/google/cloud/triggers/test_datafusion.py index 8e85bd0504121..6cc2397bafdbc 100644 --- a/tests/providers/google/cloud/triggers/test_datafusion.py +++ b/tests/providers/google/cloud/triggers/test_datafusion.py @@ -36,6 +36,7 @@ PIPELINE_NAME = "shrubberyPipeline" PIPELINE = {"test": "pipeline"} PIPELINE_ID = "test_pipeline_id" +PIPELINE_TYPE = "batch" INSTANCE_URL = "http://datafusion.instance.com" NAMESPACE = "TEST_NAMESPACE" RUNTIME_ARGS = {"arg1": "a", "arg2": "b"} @@ -50,6 +51,7 @@ def trigger(): namespace=NAMESPACE, pipeline_name=PIPELINE_NAME, pipeline_id=PIPELINE_ID, + pipeline_type=PIPELINE_TYPE, poll_interval=TEST_POLL_INTERVAL, gcp_conn_id=TEST_GCP_PROJECT_ID, ) @@ -68,6 +70,7 @@ def test_start_pipeline_trigger_serialization_should_execute_successfully(self, "namespace": NAMESPACE, "pipeline_name": PIPELINE_NAME, "pipeline_id": PIPELINE_ID, + "pipeline_type": PIPELINE_TYPE, "gcp_conn_id": TEST_GCP_PROJECT_ID, "success_states": None, } diff --git a/tests/providers/google/cloud/utils/test_datafusion.py b/tests/providers/google/cloud/utils/test_datafusion.py new file mode 100644 index 0000000000000..ab1f3a0927efe --- /dev/null +++ b/tests/providers/google/cloud/utils/test_datafusion.py @@ -0,0 +1,37 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import pytest + +from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType + + +class TestDataFusionPipelineType: + @pytest.mark.parametrize( + "str_value, expected_item", + [ + ("batch", DataFusionPipelineType.BATCH), + ("stream", DataFusionPipelineType.STREAM), + ], + ) + def test_from_str(self, str_value, expected_item): + assert DataFusionPipelineType.from_str(str_value) == expected_item + + def test_from_str_error(self): + with pytest.raises(ValueError): + DataFusionPipelineType.from_str("non-existing value")