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')
