Passing the environment variables to ContainerOp

When the DSL bridge code was written, ContainerOp did not support env, so we did not pass it. Now we're adding the passing code.
Added test that chacks that the env variables get to the ContainerOp.

sdk/python/kfp/components/_dsl_bridge.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from collections import OrderedDict
+from typing import Mapping
 from ._structures import ConcatPlaceholder, IfPlaceholder, InputValuePlaceholder, InputPathPlaceholder, IsPresentPlaceholder, OutputPathPlaceholder, TaskSpec
 from ._components import _generate_output_file_name, _default_component_name
 
@@ -123,12 +124,13 @@ def expand_argument_list(argument_list):
         command=expanded_command,
         arguments=expanded_args,
         output_paths=output_paths,
+        env=container_spec.env,
     )
 
 
 _dummy_pipeline=None
 
-def _create_container_op_from_resolved_task(name:str, container_image:str, command=None, arguments=None, output_paths=None):
+def _create_container_op_from_resolved_task(name:str, container_image:str, command=None, arguments=None, output_paths=None, env : Mapping[str, str]=None):
     from .. import dsl
     global _dummy_pipeline
     need_dummy = dsl.Pipeline._default_pipeline is None
@@ -155,7 +157,11 @@ def _create_container_op_from_resolved_task(name:str, container_image:str, comma
         arguments=arguments,
         file_outputs=output_paths_for_container_op,
     )
-
+    if env:
+        from kubernetes import client as k8s_client
+        for name, value in env.items():
+            task.add_env_variable(k8s_client.V1EnvVar(name=name, value=value))
+
     if need_dummy:
         _dummy_pipeline.__exit__()
 

sdk/python/tests/components/test_components.py

@@ -458,6 +458,24 @@ def test_command_if_input_value_then(self):
         task_else = task_factory1()
         self.assertEqual(task_else.arguments, [])
 
+    def test_handling_env(self):
+        component_text = '''\
+implementation:
+  container:
+    image: busybox
+    env:
+      key1: value 1
+      key2: value 2
+'''
+        task_factory1 = comp.load_component_from_text(component_text)
+
+        import kfp
+        with kfp.dsl.Pipeline('Dummy'): #Forcing the TaskSpec conversion to ContainerOp
+            task1 = task_factory1()
+        actual_env = {env_var.name: env_var.value for env_var in task1.env_variables}
+        expected_env = {'key1': 'value 1', 'key2': 'value 2'}
+        self.assertDictEqual(expected_env, actual_env)
+
 
 if __name__ == '__main__':
     unittest.main()