SDK/DSL: ContainerOp.add_pvolume - Fix volume passed in add_volume (#…

…2306)

Signed-off-by: Ilias Katsakioris <elikatsis@arrikto.com>

sdk/python/kfp/dsl/_container_op.py

@@ -1182,7 +1182,8 @@ def add_pvolumes(self,
                     self.dependent_names.extend(pvolume.dependent_names)
                 else:
                     pvolume = PipelineVolume(volume=pvolume)
-                self.pvolumes[mount_path] = pvolume.after(self)
+                pvolume = pvolume.after(self)
+                self.pvolumes[mount_path] = pvolume
                 self.add_volume(pvolume)
                 self._container.add_volume_mount(V1VolumeMount(
                     name=pvolume.name,

sdk/python/tests/dsl/container_op_tests.py

@@ -17,7 +17,7 @@
 from kubernetes.client.models import V1EnvVar, V1VolumeMount
 
 import kfp
-from kfp.dsl import ContainerOp, UserContainer, Sidecar
+from kfp.dsl import ContainerOp, UserContainer, Sidecar, PipelineVolume
 
 
 class TestContainerOp(unittest.TestCase):
@@ -89,3 +89,12 @@ def test_deprecation_warnings(self):
       op.add_volume_mount(V1VolumeMount(
         mount_path='/secret/gcp-credentials',
         name='gcp-credentials'))
+
+
+  def test_add_pvolumes(self):
+    pvolume = PipelineVolume(pvc='test')
+    op = ContainerOp(name='op1', image='image', pvolumes={'/mnt': pvolume})
+
+    self.assertEqual(pvolume.dependent_names, [])
+    self.assertEqual(op.pvolume.dependent_names, [op.name])
+    self.assertEqual(op.volumes[0].dependent_names, [op.name])