From ea85bfef1a44e52df34202f374613c043be7a4d9 Mon Sep 17 00:00:00 2001 From: muchida Date: Thu, 10 Oct 2019 16:33:25 -0700 Subject: [PATCH] Partially revert kubeflow_dag_runner_test to be backward compatible PiperOrigin-RevId: 274064819 --- .../kubeflow/kubeflow_dag_runner_test.py | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/tfx/orchestration/kubeflow/kubeflow_dag_runner_test.py b/tfx/orchestration/kubeflow/kubeflow_dag_runner_test.py index c070605d4b..9335a6123c 100644 --- a/tfx/orchestration/kubeflow/kubeflow_dag_runner_test.py +++ b/tfx/orchestration/kubeflow/kubeflow_dag_runner_test.py @@ -172,6 +172,17 @@ def testVolumeMountingPipelineOperatorFuncs(self): ] self.assertEqual(2, len(container_templates)) + volumes = [{ + 'name': 'my-volume-name', + 'persistentVolumeClaim': { + 'claimName': 'my-persistent-volume-claim' + } + }] + + # Check that the PVC is specified for kfp<=0.1.31.1. + if 'volumes' in pipeline['spec']: + self.assertEqual(volumes, pipeline['spec']['volumes']) + for template in container_templates: # Check that each container has the volume mounted. self.assertEqual([{ @@ -179,13 +190,9 @@ def testVolumeMountingPipelineOperatorFuncs(self): 'mountPath': '/mnt/volume-mount-path' }], template['container']['volumeMounts']) - # Check that each template has the PVC specified. - self.assertEqual([{ - 'name': 'my-volume-name', - 'persistentVolumeClaim': { - 'claimName': 'my-persistent-volume-claim' - } - }], template['volumes']) + # Check that each template has the PVC specified for kfp>=0.1.31.2. + if 'volumes' in template: + self.assertEqual(volumes, template['volumes']) if __name__ == '__main__':