fix test

components/gcp/container/component_sdk/python/kfp_component/google/dataflow/_launch_python.py

@@ -33,7 +33,8 @@ def launch_python(python_file_path, project_id, region, staging_dir=None, requir
 
     Args:
         python_file_path (str): The gcs or local path to the python file to run.
-        project_id (str): The ID of the parent project.
+        project_id (str): The ID of the GCP project to run the Dataflow job.
+        region (str): The GCP region to run the Dataflow job.
         staging_dir (str): Optional. The GCS directory for keeping staging files.
             A random subdirectory will be created under the directory to keep job info
             for resuming the job in case of failure and it will be passed as

components/gcp/container/component_sdk/python/setup.py

@@ -36,8 +36,6 @@
         'Intended Audience :: Science/Research',
         'License :: OSI Approved :: Apache Software License',
         'Programming Language :: Python',
-        'Programming Language :: Python :: 2',
-        'Programming Language :: Python :: 2.7',
         'Programming Language :: Python :: 3',
         'Programming Language :: Python :: 3.5',
         'Programming Language :: Python :: 3.6',

components/gcp/container/component_sdk/python/tests/google/dataflow/test__launch_python.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Google LLC
+# Copyright 2021 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -29,27 +29,27 @@
 @mock.patch(MODULE + '.subprocess')
 class LaunchPythonTest(unittest.TestCase):
 
-    def test_launch_python_succeed(self, mock_subprocess, mock_process, 
+    def test_launch_python_succeed(self, mock_subprocess, mock_process,
         mock_client, mock_context, mock_stage_file, mock_display, mock_storage):
         mock_context().__enter__().context_id.return_value = 'ctx-1'
         mock_storage.Client().bucket().blob().exists.return_value = False
         mock_process().read_lines.return_value = [
-            b'https://console.cloud.google.com/dataflow/locations/us-central1/jobs/job-1?project=project-1'
+            b'https://console.cloud.google.com/dataflow/jobs/us-central1/job-1?project=project-1'
         ]
         expected_job = {
             'id': 'job-1',
             'currentState': 'JOB_STATE_DONE'
         }
         mock_client().get_job.return_value = expected_job
 
-        result = launch_python('/tmp/test.py', 'project-1', staging_dir='gs://staging/dir')
+        result = launch_python('/tmp/test.py', 'project-1', 'us-central1', staging_dir='gs://staging/dir')
 
         self.assertEqual(expected_job, result)
         mock_storage.Client().bucket().blob().upload_from_string.assert_called_with(
             'job-1,us-central1'
         )
 
-    def test_launch_python_retry_succeed(self, mock_subprocess, mock_process, 
+    def test_launch_python_retry_succeed(self, mock_subprocess, mock_process,
         mock_client, mock_context, mock_stage_file, mock_display, mock_storage):
         mock_context().__enter__().context_id.return_value = 'ctx-1'
         mock_storage.Client().bucket().blob().exists.return_value = True
@@ -60,20 +60,19 @@ def test_launch_python_retry_succeed(self, mock_subprocess, mock_process,
         }
         mock_client().get_job.return_value = expected_job
 
-        result = launch_python('/tmp/test.py', 'project-1', staging_dir='gs://staging/dir')
+        result = launch_python('/tmp/test.py', 'project-1', 'us-central1', staging_dir='gs://staging/dir')
 
         self.assertEqual(expected_job, result)
         mock_process.assert_not_called()
 
-    def test_launch_python_no_job_created(self, mock_subprocess, mock_process, 
+    def test_launch_python_no_job_created(self, mock_subprocess, mock_process,
         mock_client, mock_context, mock_stage_file, mock_display, mock_storage):
         mock_context().__enter__().context_id.return_value = 'ctx-1'
         mock_process().read_lines.return_value = [
             b'no job id',
             b'no job id'
         ]
 
-        result = launch_python('/tmp/test.py', 'project-1')
+        result = launch_python('/tmp/test.py', 'project-1', 'us-central1')
 
         self.assertEqual(None, result)
-

components/gcp/container/component_sdk/python/tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = py27,py35,py36,py37,py38
+envlist = py35,py36,py37,py38
 skip_missing_interpreters = true
 
 [testenv]