def testFlexTemplatesCreateJob(self): name = 'job name' gcs_path = 'gs://my_gcs_path' template_args = apis.TemplateArguments( job_name=name, gcs_location=gcs_path, worker_region=base.DEFAULT_REGION, project_id=self.Project()) launch_params = apis.Templates.FLEX_TEMPLATE_PARAMETER( jobName=template_args.job_name, containerSpecGcsPath=template_args.gcs_location, environment=apis.Templates.FLEX_TEMPLATE_ENVIRONMENT( workerRegion=base.DEFAULT_REGION, additionalExperiments=[]), parameters=None) request = apis.Templates.LAUNCH_FLEX_TEMPLATE_REQUEST( launchParameter=launch_params) wrapper_req = (apis.GetMessagesModule(). DataflowProjectsLocationsFlexTemplatesLaunchRequest) wrapped_request = wrapper_req(projectId=self.Project(), location=base.DEFAULT_REGION, launchFlexTemplateRequest=request) self.mocked_client.projects_locations_flexTemplates.Launch.Expect( wrapped_request, apis.GetMessagesModule().LaunchFlexTemplateResponse()) apis.Templates.CreateJobFromFlexTemplate(template_args)
def testTemplatesCreate(self): name = 'job name' gcs_path = 'gs://my_gcs_path' staging_path = 'gs://my_staging_path' request = apis.Templates.CREATE_REQUEST( jobName=name, location=base.DEFAULT_REGION, gcsPath=gcs_path, environment=apis.GetMessagesModule().RuntimeEnvironment( tempLocation=staging_path)) wrapper_req = (apis.GetMessagesModule(). DataflowProjectsLocationsTemplatesCreateRequest) wrapped_request = wrapper_req(createJobFromTemplateRequest=request, projectId=self.Project(), location=base.DEFAULT_REGION) self.mocked_client.projects_locations_templates.Create.Expect( wrapped_request, apis.GetMessagesModule().Job()) template_args = apis.TemplateArguments(job_name=name, gcs_location=gcs_path, project_id=self.Project(), staging_location=staging_path) apis.Templates.Create(template_args)
def testPrivateIPTemplatesCreate(self): name = 'job name' gcs_path = 'gs://my_gcs_path' staging_path = 'gs://my_staging_path' ip_configuration_enum = apis.GetMessagesModule( ).RuntimeEnvironment.IpConfigurationValueValuesEnum ip_private = ip_configuration_enum.WORKER_IP_PRIVATE request = apis.Templates.CREATE_REQUEST( jobName=name, location=base.DEFAULT_REGION, gcsPath=gcs_path, environment=apis.GetMessagesModule().RuntimeEnvironment( tempLocation=staging_path, ipConfiguration=ip_private)) wrapper_req = (apis.GetMessagesModule(). DataflowProjectsLocationsTemplatesCreateRequest) wrapped_request = wrapper_req(createJobFromTemplateRequest=request, projectId=self.Project(), location=base.DEFAULT_REGION) self.mocked_client.projects_locations_templates.Create.Expect( wrapped_request, apis.GetMessagesModule().Job()) template_args = apis.TemplateArguments(job_name=name, gcs_location=gcs_path, project_id=self.Project(), staging_location=staging_path, disable_public_ips=True) apis.Templates.Create(template_args)
def testLaunchDynamicTemplate(self): name = 'myjob' gcs_path = 'gs://dynamic_template_path' staging_path = 'gs://my_staging_path' parameters = apis.Templates.LAUNCH_TEMPLATE_PARAMETERS( jobName=name, environment=apis.GetMessagesModule().RuntimeEnvironment( tempLocation=staging_path), update=False) wrapper_req = (apis.GetMessagesModule(). DataflowProjectsLocationsTemplatesLaunchRequest) wrapped_request = wrapper_req( dynamicTemplate_gcsPath=gcs_path, dynamicTemplate_stagingLocation=staging_path, launchTemplateParameters=parameters, projectId=self.Project(), location=base.DEFAULT_REGION, validateOnly=False) self.mocked_client.projects_locations_templates.Launch.Expect( wrapped_request, apis.GetMessagesModule().LaunchTemplateResponse()) template_args = apis.TemplateArguments(job_name=name, gcs_location=gcs_path, project_id=self.Project(), staging_location=staging_path) apis.Templates.LaunchDynamicTemplate(template_args)
def _StatusArgToFilter(self, status, region=None): """Return a string describing the job status. Args: status: The job status enum region: The region argument, to select the correct wrapper message. Returns: string describing the job status """ filter_value_enum = None if region: filter_value_enum = ( apis.GetMessagesModule() .DataflowProjectsLocationsJobsListRequest.FilterValueValuesEnum) else: filter_value_enum = ( apis.GetMessagesModule() .DataflowProjectsJobsAggregatedRequest.FilterValueValuesEnum) value_map = { 'all': filter_value_enum.ALL, 'terminated': filter_value_enum.TERMINATED, 'active': filter_value_enum.ACTIVE, } return value_map.get(status, filter_value_enum.ALL)
def SetUp(self): state_enum = apis.GetMessagesModule().Job.CurrentStateValueValuesEnum self.terminated_states = [state_enum.JOB_STATE_CANCELLED, state_enum.JOB_STATE_DONE, state_enum.JOB_STATE_FAILED, state_enum.JOB_STATE_STOPPED, state_enum.JOB_STATE_UPDATED] type_enum = apis.GetMessagesModule().Job.TypeValueValuesEnum self.valid_types = [type_enum.JOB_TYPE_BATCH, type_enum.JOB_TYPE_STREAMING]
def testMessagesList(self): start_time = 'start_time' end_time = 'end_time' page_size = 1337 page_token = 'page_token' minimum_importance = ( apis.Messages.LIST_REQUEST.MinimumImportanceValueValuesEnum. JOB_MESSAGE_DETAILED) request = apis.Messages.LIST_REQUEST( jobId=base.JOB_1_ID, projectId=self.Project(), location=base.DEFAULT_REGION, minimumImportance=minimum_importance, startTime=start_time, endTime=end_time, pageSize=page_size, pageToken=page_token) self.mocked_client.projects_locations_jobs_messages.List.Expect( request, apis.GetMessagesModule().ListJobMessagesResponse()) apis.Messages.List(job_id=base.JOB_1_ID, project_id=self.Project(), minimum_importance=minimum_importance, start_time=start_time, end_time=end_time, page_size=page_size, page_token=page_token)
def testJobsDrain(self): job_id = base.JOB_1_ID project_id = self.Project() state = (apis.GetMessagesModule().Job.RequestedStateValueValuesEnum. JOB_STATE_DRAINED) job = apis.GetMessagesModule().Job(requestedState=state) request = apis.Jobs.UPDATE_REQUEST(jobId=job_id, projectId=project_id, location=base.DEFAULT_REGION, job=job) self.mocked_client.projects_locations_jobs.Update.Expect( request, apis.GetMessagesModule().Job()) apis.Jobs.Drain(job_id=job_id, project_id=project_id)
def testMetricsGet(self): request = apis.Metrics.GET_REQUEST(jobId=base.JOB_1_ID, projectId=self.Project(), location=base.DEFAULT_REGION, startTime='start_time') self.mocked_client.projects_locations_jobs.GetMetrics.Expect( request, apis.GetMessagesModule().JobMetrics()) apis.Metrics.Get(job_id=base.JOB_1_ID, project_id=self.Project(), start_time='start_time')
def testJobsGet(self): request = apis.Jobs.GET_REQUEST( jobId=base.JOB_1_ID, projectId=self.Project(), location=base.DEFAULT_REGION, view=apis.Jobs.GET_REQUEST.ViewValueValuesEnum.JOB_VIEW_ALL) self.mocked_client.projects_locations_jobs.Get.Expect( request, apis.GetMessagesModule().Job()) apis.Jobs.Get( job_id=base.JOB_1_ID, project_id=self.Project(), view=apis.Jobs.GET_REQUEST.ViewValueValuesEnum.JOB_VIEW_ALL)
def _JobTypeForJob(job_type): """Return a string describing the job type. Args: job_type: The job type enum Returns: string describing the job type """ type_value_enum = apis.GetMessagesModule().Job.TypeValueValuesEnum value_map = { type_value_enum.JOB_TYPE_BATCH: 'Batch', type_value_enum.JOB_TYPE_STREAMING: 'Streaming', } return value_map.get(job_type, 'Unknown')
def testLaunchDynamicTemplate_httpError(self): name = 'myjob' gcs_path = 'gs://dynamic_template_path' staging_path = 'gs://my_staging_path' parameters = apis.Templates.LAUNCH_TEMPLATE_PARAMETERS( jobName=name, environment=apis.GetMessagesModule().RuntimeEnvironment( tempLocation=staging_path), update=False) wrapper_req = (apis.GetMessagesModule(). DataflowProjectsLocationsTemplatesLaunchRequest) wrapped_request = wrapper_req( dynamicTemplate_gcsPath=gcs_path, dynamicTemplate_stagingLocation=staging_path, launchTemplateParameters=parameters, projectId=self.Project(), location=base.DEFAULT_REGION, validateOnly=False) self.mocked_client.projects_locations_templates.Launch.Expect( wrapped_request, exception=apitools_exceptions.HttpError.FromResponse( http_wrapper.Response( info={'status': 400}, content='{"error": ' '{"message": "The workflow could not be created."}}', request_url= 'https://dataflow.googleapis.com/v1b3/request-url'))) template_args = apis.TemplateArguments(job_name=name, gcs_location=gcs_path, project_id=self.Project(), staging_location=staging_path) with self.AssertRaisesExceptionMatches( exceptions.HttpException, 'The workflow could not be created.'): apis.Templates.LaunchDynamicTemplate(template_args)
def _StatusArgToFilter(self, status): """Return a string describing the job status. Args: status: The job status enum Returns: string describing the job status """ filter_value_enum = (apis.GetMessagesModule( ).DataflowProjectsJobsListRequest.FilterValueValuesEnum) value_map = { 'all': filter_value_enum.ALL, 'terminated': filter_value_enum.TERMINATED, 'active': filter_value_enum.ACTIVE, } return value_map.get(status, filter_value_enum.ALL)
def _StatusForJob(job_state): """Return a string describing the job state. Args: job_state: The job state enum Returns: string describing the job state """ state_value_enum = apis.GetMessagesModule().Job.CurrentStateValueValuesEnum value_map = { state_value_enum.JOB_STATE_CANCELLED: 'Cancelled', state_value_enum.JOB_STATE_DONE: 'Done', state_value_enum.JOB_STATE_FAILED: 'Failed', state_value_enum.JOB_STATE_RUNNING: 'Running', state_value_enum.JOB_STATE_STOPPED: 'Stopped', state_value_enum.JOB_STATE_UPDATED: 'Updated', } return value_map.get(job_state, 'Unknown')