def testMainWithV1Producer(self):
"""Tests the entrypoint with data passing with conventional KFP components.
This test case emulates the following scenario:
- User provides a function, namely `test_func`.
- In test function, there are an input parameter (`test_param`) and an input
artifact (`test_artifact`). And the user code generates an output
artifact (`test_output1`) and an output parameter (`test_output2`).
- The specified metadata JSON file location is at
'executor_output_metadata.json'
- The inputs of this step are all provided by conventional KFP components.
"""
# Set mocked user function.
self._import_func.return_value = main.test_func
entrypoint.main(
executor_metadata_json_file=_OUTPUT_METADATA_JSON_LOCATION,
function_name='test_func',
test_param_input_argo_param='hello from producer',
test_artifact_input_path='gs://root/producer/output',
test_output1_artifact_output_path='gs://root/consumer/output1',
test_output2_parameter_output_path='gs://root/consumer/output2')
self._mock_gcs_write.assert_called_with(
path=_OUTPUT_METADATA_JSON_LOCATION,
content=_EXPECTED_EXECUTOR_OUTPUT_1)
def testMainWithV2Producer(self):
"""Tests the entrypoint with data passing with new-styled KFP components.
This test case emulates a similar scenario as testMainWithV1Producer, except
for that the inputs of this step are all provided by a new-styled KFP
component.
"""
# Set mocked user function.
self._import_func.return_value = main.test_func2
# Set GFile read function
self._mock_gcs_read.return_value = _PRODUCER_EXECUTOR_OUTPUT
entrypoint.main(
executor_metadata_json_file=_OUTPUT_METADATA_JSON_LOCATION,
function_name='test_func2',
test_param_input_param_metadata_file=
'gs://root/producer/executor_output_metadata.json',
test_param_input_field_name='param_output',
test_artifact_input_artifact_metadata_file=
'gs://root/producer/executor_output_metadata.json',
test_artifact_input_output_name='artifact_output',
test_output1_artifact_output_path='gs://root/consumer/output1',
test_output2_parameter_output_path='gs://root/consumer/output2')
self._mock_gcs_write.assert_called_with(
path=_OUTPUT_METADATA_JSON_LOCATION,
content=_EXPECTED_EXECUTOR_OUTPUT_1)
def testMainWithV2Producer(self):
"""Tests the entrypoint with data passing with new-styled KFP components.
This test case emulates a similar scenario as testMainWithV1Producer, except
for that the inputs of this step are all provided by a new-styled KFP
component.
"""
# Set mocked user function.
self._import_func.return_value = main.test_func2
entrypoint.main(executor_input_str=_TEST_EXECUTOR_INPUT_V2_PRODUCER,
function_name='test_func2',
output_metadata_path=_OUTPUT_METADATA_JSON_LOCATION)
# Check the actual executor output.
with open(_OUTPUT_METADATA_JSON_LOCATION, 'r') as f:
self.assertEqual(f.read(), _EXPECTED_EXECUTOR_OUTPUT)
def testMainWithV1Producer(self):
"""Tests the entrypoint with data passing with conventional KFP components.
This test case emulates the following scenario:
- User provides a function, namely `test_func`.
- In test function, there are an input parameter (`test_param`) and an input
artifact (`test_artifact`). And the user code generates an output
artifact (`test_output1`) and an output parameter (`test_output2`).
- The specified metadata JSON file location is at
'executor_output_metadata.json'
- The inputs of this step are all provided by conventional KFP components.
"""
# Set mocked user function.
self._import_func.return_value = main.test_func
entrypoint.main(executor_input_str=_TEST_EXECUTOR_INPUT_V1_PRODUCER,
function_name='test_func',
output_metadata_path=_OUTPUT_METADATA_JSON_LOCATION)
# Check the actual executor output.
with open(_OUTPUT_METADATA_JSON_LOCATION, 'r') as f:
self.assertEqual(f.read(), _EXPECTED_EXECUTOR_OUTPUT)