def __init__(
self,
*,
job_name: str,
job_definition: str,
job_queue: str,
overrides: dict,
array_properties: Optional[dict] = None,
parameters: Optional[dict] = None,
job_id: Optional[str] = None,
waiters: Optional[Any] = None,
max_retries: Optional[int] = None,
status_retries: Optional[int] = None,
aws_conn_id: Optional[str] = None,
region_name: Optional[str] = None,
tags: Optional[dict] = None,
**kwargs,
): # pylint: disable=too-many-arguments
BaseOperator.__init__(self, **kwargs)
self.job_id = job_id
self.job_name = job_name
self.job_definition = job_definition
self.job_queue = job_queue
self.overrides = overrides or {}
self.array_properties = array_properties or {}
self.parameters = parameters or {}
self.waiters = waiters
self.tags = tags or {}
self.hook = AwsBatchClientHook(
max_retries=max_retries,
status_retries=status_retries,
aws_conn_id=aws_conn_id,
region_name=region_name,
)
def __init__(
self,
job_name,
job_definition,
job_queue,
overrides,
array_properties=None,
parameters=None,
job_id=None,
waiters=None,
max_retries=None,
status_retries=None,
aws_conn_id=None,
region_name=None,
**kwargs,
): # pylint: disable=too-many-arguments
BaseOperator.__init__(self, **kwargs)
self.job_id = job_id
self.job_name = job_name
self.job_definition = job_definition
self.job_queue = job_queue
self.overrides = overrides
self.array_properties = array_properties or {}
self.parameters = parameters
self.waiters = waiters
self.hook = AwsBatchClientHook(
max_retries=max_retries,
status_retries=status_retries,
aws_conn_id=aws_conn_id,
region_name=region_name,
)
def get_hook(self) -> AwsBatchClientHook:
"""Create and return a AwsBatchClientHook"""
if self.hook:
return self.hook
self.hook = AwsBatchClientHook(
aws_conn_id=self.aws_conn_id,
region_name=self.region_name,
)
return self.hook
def setUp(self, get_client_type_mock):
self.get_client_type_mock = get_client_type_mock
self.batch_client = AwsBatchClientHook(
max_retries=self.MAX_RETRIES,
status_retries=self.STATUS_RETRIES,
aws_conn_id='airflow_test',
region_name=AWS_REGION,
)
self.client_mock = get_client_type_mock.return_value
assert self.batch_client.client == self.client_mock # setup client property
# don't pause in these unit tests
self.mock_delay = mock.Mock(return_value=None)
self.batch_client.delay = self.mock_delay
self.mock_exponential_delay = mock.Mock(return_value=0)
self.batch_client.exponential_delay = self.mock_exponential_delay
def setUp(self):
self.batch_client = AwsBatchClientHook(aws_conn_id='airflow_test',
region_name=AWS_REGION)