Ejemplo n.º 1
0
 def test_update_queue(self, mock_hook):
     mock_hook.return_value.update_queue.return_value = {}
     operator = CloudTasksQueueUpdateOperator(
         task_queue=Queue(name=FULL_QUEUE_PATH), task_id="id")
     operator.execute(context=None)
     mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
     mock_hook.return_value.update_queue.assert_called_once_with(
         task_queue=Queue(name=FULL_QUEUE_PATH),
         project_id=None,
         location=None,
         queue_name=None,
         update_mask=None,
         retry=None,
         timeout=None,
         metadata=None,
     )
Ejemplo n.º 2
0
    def update_queue(
            self,
            task_queue: Queue,
            project_id: str,
            location: Optional[str] = None,
            queue_name: Optional[str] = None,
            update_mask: Optional[FieldMask] = None,
            retry: Optional[Retry] = None,
            timeout: Optional[float] = None,
            metadata: Optional[Sequence[Tuple[str, str]]] = None) -> Queue:
        """
        Updates a queue in Cloud Tasks.

        :param task_queue: The task queue to update.
            This method creates the queue if it does not exist and updates the queue if
            it does exist. The queue's name must be specified.
        :type task_queue: dict or class google.cloud.tasks_v2.types.Queue
        :param project_id: (Optional) The ID of the  GCP project that owns the Cloud Tasks.
            If set to None or missing, the default project_id from the GCP connection is used.
        :type project_id: str
        :param location: (Optional) The location name in which the queue will be updated.
            If provided, it will be used to construct the full queue path.
        :type location: str
        :param queue_name: (Optional) The queue's name.
            If provided, it will be used to construct the full queue path.
        :type queue_name: str
        :param update_mask: A mast used to specify which fields of the queue are being updated.
            If empty, then all fields will be updated.
            If a dict is provided, it must be of the same form as the protobuf message.
        :type update_mask: dict or class google.cloud.tasks_v2.types.FieldMask
        :param retry: (Optional) A retry object used to retry requests.
            If None is specified, requests will not be retried.
        :type retry: google.api_core.retry.Retry
        :param timeout: (Optional) The amount of time, in seconds, to wait for the request
            to complete. Note that if retry is specified, the timeout applies to each
            individual attempt.
        :type timeout: float
        :param metadata: (Optional) Additional metadata that is provided to the method.
        :type metadata: sequence[tuple[str, str]]]
        :rtype: google.cloud.tasks_v2.types.Queue
        """

        client = self.get_conn()

        if queue_name and location:
            full_queue_name = CloudTasksClient.queue_path(
                project_id, location, queue_name)
            if isinstance(task_queue, Queue):
                task_queue.name = full_queue_name
            elif isinstance(task_queue, dict):
                task_queue['name'] = full_queue_name
            else:
                raise AirflowException('Unable to set queue_name.')
        return client.update_queue(
            queue=task_queue,
            update_mask=update_mask,
            retry=retry,
            timeout=timeout,
            metadata=metadata,
        )
Ejemplo n.º 3
0
    def execute(self, context):
        hook = CloudTasksHook(
            gcp_conn_id=self.gcp_conn_id,
            impersonation_chain=self.impersonation_chain,
        )
        try:
            queue = hook.create_queue(
                location=self.location,
                task_queue=self.task_queue,
                project_id=self.project_id,
                queue_name=self.queue_name,
                retry=self.retry,
                timeout=self.timeout,
                metadata=self.metadata,
            )
        except AlreadyExists:
            queue = hook.get_queue(
                location=self.location,
                project_id=self.project_id,
                queue_name=self.queue_name,
                retry=self.retry,
                timeout=self.timeout,
                metadata=self.metadata,
            )

        return Queue.to_dict(queue)
Ejemplo n.º 4
0
 def test_create_queue(self, mock_hook):
     mock_hook.return_value.create_queue.return_value = {}
     operator = CloudTasksQueueCreateOperator(location=LOCATION,
                                              task_queue=Queue(),
                                              task_id="id")
     operator.execute(context=None)
     mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
     mock_hook.return_value.create_queue.assert_called_once_with(
         location=LOCATION,
         task_queue=Queue(),
         project_id=None,
         queue_name=None,
         retry=None,
         timeout=None,
         metadata=None,
     )
Ejemplo n.º 5
0
    def execute(self, context: 'Context'):
        hook = CloudTasksHook(
            gcp_conn_id=self.gcp_conn_id,
            impersonation_chain=self.impersonation_chain,
        )
        try:
            queue = hook.create_queue(
                location=self.location,
                task_queue=self.task_queue,
                project_id=self.project_id,
                queue_name=self.queue_name,
                retry=self.retry,
                timeout=self.timeout,
                metadata=self.metadata,
            )
        except AlreadyExists:
            if self.queue_name is None:
                raise RuntimeError("The queue name should be set here!")
            queue = hook.get_queue(
                location=self.location,
                project_id=self.project_id,
                queue_name=self.queue_name,
                retry=self.retry,
                timeout=self.timeout,
                metadata=self.metadata,
            )

        return Queue.to_dict(queue)
Ejemplo n.º 6
0
    def test_update_queue(self, get_conn):
        result = self.hook.update_queue(
            task_queue=Queue(state=3),
            location=LOCATION,
            queue_name=QUEUE_ID,
            project_id=PROJECT_ID,
        )

        self.assertIs(result, API_RESPONSE)

        get_conn.return_value.update_queue.assert_called_once_with(
            queue=Queue(name=FULL_QUEUE_PATH, state=3),
            update_mask=None,
            retry=None,
            timeout=None,
            metadata=None,
        )
Ejemplo n.º 7
0
    def test_create_queue(self, get_conn):
        result = self.hook.create_queue(
            location=LOCATION,
            task_queue=Queue(),
            queue_name=QUEUE_ID,
            project_id=PROJECT_ID,
        )

        self.assertIs(result, API_RESPONSE)

        get_conn.return_value.create_queue.assert_called_once_with(
            parent=FULL_LOCATION_PATH,
            queue=Queue(name=FULL_QUEUE_PATH),
            retry=None,
            timeout=None,
            metadata=None,
        )
Ejemplo n.º 8
0
    def update_queue(
            self,
            task_queue: Queue,
            project_id: str = PROVIDE_PROJECT_ID,
            location: Optional[str] = None,
            queue_name: Optional[str] = None,
            update_mask: Optional[FieldMask] = None,
            retry: Union[Retry, _MethodDefault] = DEFAULT,
            timeout: Optional[float] = None,
            metadata: Sequence[Tuple[str, str]] = (),
    ) -> Queue:
        """
        Updates a queue in Cloud Tasks.

        :param task_queue: The task queue to update.
            This method creates the queue if it does not exist and updates the queue if
            it does exist. The queue's name must be specified.
        :param project_id: (Optional) The ID of the Google Cloud project that owns the Cloud Tasks.
            If set to None or missing, the default project_id from the Google Cloud connection is used.
        :param location: (Optional) The location name in which the queue will be updated.
            If provided, it will be used to construct the full queue path.
        :param queue_name: (Optional) The queue's name.
            If provided, it will be used to construct the full queue path.
        :param update_mask: A mast used to specify which fields of the queue are being updated.
            If empty, then all fields will be updated.
            If a dict is provided, it must be of the same form as the protobuf message.
        :param retry: (Optional) A retry object used to retry requests.
            If None is specified, requests will not be retried.
        :param timeout: (Optional) The amount of time, in seconds, to wait for the request
            to complete. Note that if retry is specified, the timeout applies to each
            individual attempt.
        :param metadata: (Optional) Additional metadata that is provided to the method.
        :rtype: google.cloud.tasks_v2.types.Queue
        """
        client = self.get_conn()

        if queue_name and location:
            full_queue_name = f"projects/{project_id}/locations/{location}/queues/{queue_name}"
            if isinstance(task_queue, Queue):
                task_queue.name = full_queue_name
            elif isinstance(task_queue, dict):
                task_queue['name'] = full_queue_name
            else:
                raise AirflowException('Unable to set queue_name.')
        return client.update_queue(
            request={
                'queue': task_queue,
                'update_mask': update_mask
            },
            retry=retry,
            timeout=timeout,
            metadata=metadata,
        )
Ejemplo n.º 9
0
    def test_resume_queue(self, get_conn):
        result = self.hook.resume_queue(location=LOCATION,
                                        queue_name=QUEUE_ID,
                                        project_id=PROJECT_ID)

        self.assertEqual(result, Queue(name=FULL_QUEUE_PATH))

        get_conn.return_value.resume_queue.assert_called_once_with(
            request=dict(name=FULL_QUEUE_PATH),
            retry=None,
            timeout=None,
            metadata=())
Ejemplo n.º 10
0
    def test_update_queue(self, mock_hook):
        mock_hook.return_value.update_queue.return_value = TEST_QUEUE
        operator = CloudTasksQueueUpdateOperator(
            task_queue=Queue(name=FULL_QUEUE_PATH), task_id="id")

        result = operator.execute(context=None)

        self.assertEqual({'name': FULL_QUEUE_PATH, 'state': 0}, result)
        mock_hook.assert_called_once_with(
            gcp_conn_id=GCP_CONN_ID,
            impersonation_chain=None,
        )
        mock_hook.return_value.update_queue.assert_called_once_with(
            task_queue=Queue(name=FULL_QUEUE_PATH),
            project_id=None,
            location=None,
            queue_name=None,
            update_mask=None,
            retry=None,
            timeout=None,
            metadata=None,
        )
Ejemplo n.º 11
0
 def execute(self, context: 'Context'):
     hook = CloudTasksHook(
         gcp_conn_id=self.gcp_conn_id,
         impersonation_chain=self.impersonation_chain,
     )
     queue = hook.resume_queue(
         location=self.location,
         queue_name=self.queue_name,
         project_id=self.project_id,
         retry=self.retry,
         timeout=self.timeout,
         metadata=self.metadata,
     )
     return Queue.to_dict(queue)
Ejemplo n.º 12
0
    def test_list_queues(self, get_conn):
        result = self.hook.list_queues(location=LOCATION,
                                       project_id=PROJECT_ID)

        self.assertEqual(result, [Queue(name=FULL_QUEUE_PATH)])

        get_conn.return_value.list_queues.assert_called_once_with(
            request=dict(parent=FULL_LOCATION_PATH,
                         filter=None,
                         page_size=None),
            retry=None,
            timeout=None,
            metadata=(),
        )
Ejemplo n.º 13
0
 def execute(self, context: 'Context'):
     hook = CloudTasksHook(
         gcp_conn_id=self.gcp_conn_id,
         impersonation_chain=self.impersonation_chain,
     )
     queues = hook.list_queues(
         location=self.location,
         project_id=self.project_id,
         results_filter=self.results_filter,
         page_size=self.page_size,
         retry=self.retry,
         timeout=self.timeout,
         metadata=self.metadata,
     )
     return [Queue.to_dict(q) for q in queues]
Ejemplo n.º 14
0
def ensure_required_queues_exist():
    """
        Reads settings.CLOUDS_TASKS_QUEUES
        and creates or updates the specified queues
    """
    client = get_cloud_tasks_client()
    parent_path = cloud_tasks_parent_path()

    for queue in getattr(settings, "CLOUD_TASKS_QUEUES", []):
        queue_name = queue["name"]

        # In our task settings we expect that the queue name will not
        # include the path, otherwise moving the app, changing location
        # etc. involves changing a load of settings.
        assert ("/" not in queue_name)  # Don't specify the full path

        update_mask = ["name"]
        queue_dict = {}
        queue_dict["name"] = "%s/queues/%s" % (parent_path, queue_name)
        queue_dict["rate_limits"] = {}
        queue_dict["retry_config"] = {}

        if "rate_per_second" in queue:
            update_mask.append("rate_limits.max_dispatches_per_second")
            queue_dict["rate_limits"]["max_dispatches_per_second"] = queue[
                "rate_per_second"]

        if "rate_max_concurrent" in queue:
            update_mask.append("rate_limits.max_concurrent_dispatches")
            queue_dict["rate_limits"]["max_concurrent_dispatches"] = queue[
                "rate_max_concurrent"]

        if "retry_max_attempts" in queue:
            update_mask.append("retry_config.max_attempts")
            queue_dict["retry_config"]["max_attempts"] = queue[
                "retry_max_attempts"]

        logging.debug("Ensuring task queue is up-to-date: %s",
                      queue_dict["name"])

        from google.cloud.tasks_v2.types import Queue
        queue = Queue(**queue_dict)

        client.update_queue(
            queue=queue,
            update_mask=field_mask_pb2.FieldMask(paths=update_mask))
Ejemplo n.º 15
0
 def execute(self, context: 'Context'):
     hook = CloudTasksHook(
         gcp_conn_id=self.gcp_conn_id,
         impersonation_chain=self.impersonation_chain,
     )
     queue = hook.pause_queue(
         location=self.location,
         queue_name=self.queue_name,
         project_id=self.project_id,
         retry=self.retry,
         timeout=self.timeout,
         metadata=self.metadata,
     )
     CloudTasksQueueLink.persist(
         operator_instance=self,
         context=context,
         queue_name=queue.name,
     )
     return Queue.to_dict(queue)
Ejemplo n.º 16
0
TASK = {
    "app_engine_http_request": {  # Specify the type of request.
        "http_method": "POST",
        "relative_uri": "/example_task_handler",
        "body": "Hello".encode(),
    },
    "schedule_time": timestamp,
}

with DAG("example_gcp_tasks",
         default_args=default_args,
         schedule_interval=None) as dag:

    create_queue = CloudTasksQueueCreateOperator(
        location=LOCATION,
        task_queue=Queue(),
        queue_name=QUEUE_ID,
        retry=Retry(maximum=10.0),
        timeout=5,
        task_id="create_queue",
    )

    create_task_to_run = CloudTasksTaskCreateOperator(
        location=LOCATION,
        queue_name=QUEUE_ID,
        task=TASK,
        task_name=TASK_NAME,
        retry=Retry(maximum=10.0),
        timeout=5,
        task_id="create_task_to_run",
    )
Ejemplo n.º 17
0
class TestCloudTasksHook(unittest.TestCase):
    def setUp(self):
        with mock.patch(
                "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__",
                new=mock_base_gcp_hook_no_default_project_id,
        ):
            self.hook = CloudTasksHook(gcp_conn_id="test")

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.client_info",
        new_callable=mock.PropertyMock,
    )
    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook._get_credentials"
    )
    @mock.patch("airflow.providers.google.cloud.hooks.tasks.CloudTasksClient")
    def test_cloud_tasks_client_creation(self, mock_client, mock_get_creds,
                                         mock_client_info):
        result = self.hook.get_conn()
        mock_client.assert_called_once_with(
            credentials=mock_get_creds.return_value,
            client_info=mock_client_info.return_value)
        self.assertEqual(mock_client.return_value, result)
        self.assertEqual(self.hook._client, result)

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{"return_value.create_queue.return_value":
           API_RESPONSE},  # type: ignore
    )
    def test_create_queue(self, get_conn):
        result = self.hook.create_queue(
            location=LOCATION,
            task_queue=Queue(),
            queue_name=QUEUE_ID,
            project_id=PROJECT_ID,
        )

        self.assertIs(result, API_RESPONSE)

        get_conn.return_value.create_queue.assert_called_once_with(
            request=dict(parent=FULL_LOCATION_PATH,
                         queue=Queue(name=FULL_QUEUE_PATH)),
            retry=None,
            timeout=None,
            metadata=(),
        )

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{"return_value.update_queue.return_value":
           API_RESPONSE},  # type: ignore
    )
    def test_update_queue(self, get_conn):
        result = self.hook.update_queue(
            task_queue=Queue(state=3),
            location=LOCATION,
            queue_name=QUEUE_ID,
            project_id=PROJECT_ID,
        )

        self.assertIs(result, API_RESPONSE)

        get_conn.return_value.update_queue.assert_called_once_with(
            request=dict(queue=Queue(name=FULL_QUEUE_PATH, state=3),
                         update_mask=None),
            retry=None,
            timeout=None,
            metadata=(),
        )

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{"return_value.get_queue.return_value":
           API_RESPONSE},  # type: ignore
    )
    def test_get_queue(self, get_conn):
        result = self.hook.get_queue(location=LOCATION,
                                     queue_name=QUEUE_ID,
                                     project_id=PROJECT_ID)

        self.assertIs(result, API_RESPONSE)

        get_conn.return_value.get_queue.assert_called_once_with(
            request=dict(name=FULL_QUEUE_PATH),
            retry=None,
            timeout=None,
            metadata=())

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{
            "return_value.list_queues.return_value":
            [Queue(name=FULL_QUEUE_PATH)]
        },  # type: ignore
    )
    def test_list_queues(self, get_conn):
        result = self.hook.list_queues(location=LOCATION,
                                       project_id=PROJECT_ID)

        self.assertEqual(result, [Queue(name=FULL_QUEUE_PATH)])

        get_conn.return_value.list_queues.assert_called_once_with(
            request=dict(parent=FULL_LOCATION_PATH,
                         filter=None,
                         page_size=None),
            retry=None,
            timeout=None,
            metadata=(),
        )

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{"return_value.delete_queue.return_value": None},  # type: ignore
    )
    def test_delete_queue(self, get_conn):
        result = self.hook.delete_queue(location=LOCATION,
                                        queue_name=QUEUE_ID,
                                        project_id=PROJECT_ID)

        self.assertEqual(result, None)

        get_conn.return_value.delete_queue.assert_called_once_with(
            request=dict(name=FULL_QUEUE_PATH),
            retry=None,
            timeout=None,
            metadata=())

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{
            "return_value.purge_queue.return_value":
            Queue(name=FULL_QUEUE_PATH)
        },  # type: ignore
    )
    def test_purge_queue(self, get_conn):
        result = self.hook.purge_queue(location=LOCATION,
                                       queue_name=QUEUE_ID,
                                       project_id=PROJECT_ID)

        self.assertEqual(result, Queue(name=FULL_QUEUE_PATH))

        get_conn.return_value.purge_queue.assert_called_once_with(
            request=dict(name=FULL_QUEUE_PATH),
            retry=None,
            timeout=None,
            metadata=())

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{
            "return_value.pause_queue.return_value":
            Queue(name=FULL_QUEUE_PATH)
        },  # type: ignore
    )
    def test_pause_queue(self, get_conn):
        result = self.hook.pause_queue(location=LOCATION,
                                       queue_name=QUEUE_ID,
                                       project_id=PROJECT_ID)

        self.assertEqual(result, Queue(name=FULL_QUEUE_PATH))

        get_conn.return_value.pause_queue.assert_called_once_with(
            request=dict(name=FULL_QUEUE_PATH),
            retry=None,
            timeout=None,
            metadata=())

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{
            "return_value.resume_queue.return_value":
            Queue(name=FULL_QUEUE_PATH)
        },  # type: ignore
    )
    def test_resume_queue(self, get_conn):
        result = self.hook.resume_queue(location=LOCATION,
                                        queue_name=QUEUE_ID,
                                        project_id=PROJECT_ID)

        self.assertEqual(result, Queue(name=FULL_QUEUE_PATH))

        get_conn.return_value.resume_queue.assert_called_once_with(
            request=dict(name=FULL_QUEUE_PATH),
            retry=None,
            timeout=None,
            metadata=())

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{"return_value.create_task.return_value":
           Task(name=FULL_TASK_PATH)},  # type: ignore
    )
    def test_create_task(self, get_conn):
        result = self.hook.create_task(
            location=LOCATION,
            queue_name=QUEUE_ID,
            task=Task(),
            project_id=PROJECT_ID,
            task_name=TASK_NAME,
        )

        self.assertEqual(result, Task(name=FULL_TASK_PATH))

        get_conn.return_value.create_task.assert_called_once_with(
            request=dict(parent=FULL_QUEUE_PATH,
                         task=Task(name=FULL_TASK_PATH),
                         response_view=None),
            retry=None,
            timeout=None,
            metadata=(),
        )

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{"return_value.get_task.return_value":
           Task(name=FULL_TASK_PATH)},  # type: ignore
    )
    def test_get_task(self, get_conn):
        result = self.hook.get_task(
            location=LOCATION,
            queue_name=QUEUE_ID,
            task_name=TASK_NAME,
            project_id=PROJECT_ID,
        )

        self.assertEqual(result, Task(name=FULL_TASK_PATH))

        get_conn.return_value.get_task.assert_called_once_with(
            request=dict(name=FULL_TASK_PATH, response_view=None),
            retry=None,
            timeout=None,
            metadata=(),
        )

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{
            "return_value.list_tasks.return_value":
            [Task(name=FULL_TASK_PATH)]
        },  # type: ignore
    )
    def test_list_tasks(self, get_conn):
        result = self.hook.list_tasks(location=LOCATION,
                                      queue_name=QUEUE_ID,
                                      project_id=PROJECT_ID)

        self.assertEqual(result, [Task(name=FULL_TASK_PATH)])

        get_conn.return_value.list_tasks.assert_called_once_with(
            request=dict(parent=FULL_QUEUE_PATH,
                         response_view=None,
                         page_size=None),
            retry=None,
            timeout=None,
            metadata=(),
        )

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{"return_value.delete_task.return_value": None},  # type: ignore
    )
    def test_delete_task(self, get_conn):
        result = self.hook.delete_task(
            location=LOCATION,
            queue_name=QUEUE_ID,
            task_name=TASK_NAME,
            project_id=PROJECT_ID,
        )

        self.assertEqual(result, None)

        get_conn.return_value.delete_task.assert_called_once_with(
            request=dict(name=FULL_TASK_PATH),
            retry=None,
            timeout=None,
            metadata=())

    @mock.patch(
        "airflow.providers.google.cloud.hooks.tasks.CloudTasksHook.get_conn",
        **{"return_value.run_task.return_value":
           Task(name=FULL_TASK_PATH)},  # type: ignore
    )
    def test_run_task(self, get_conn):
        result = self.hook.run_task(
            location=LOCATION,
            queue_name=QUEUE_ID,
            task_name=TASK_NAME,
            project_id=PROJECT_ID,
        )

        self.assertEqual(result, Task(name=FULL_TASK_PATH))

        get_conn.return_value.run_task.assert_called_once_with(
            request=dict(name=FULL_TASK_PATH, response_view=None),
            retry=None,
            timeout=None,
            metadata=(),
        )
Ejemplo n.º 18
0
        "body": b"Hello",
    },
    "schedule_time": timestamp,
}

with models.DAG(
        "example_gcp_tasks",
        schedule_interval=None,  # Override to match your needs
        start_date=days_ago(1),
        tags=['example'],
) as dag:

    # Queue operations
    create_queue = CloudTasksQueueCreateOperator(
        location=LOCATION,
        task_queue=Queue(stackdriver_logging_config=dict(sampling_ratio=0.5)),
        queue_name=QUEUE_ID,
        retry=Retry(maximum=10.0),
        timeout=5,
        task_id="create_queue",
    )

    delete_queue = CloudTasksQueueDeleteOperator(
        location=LOCATION,
        queue_name=QUEUE_ID,
        task_id="delete_queue",
    )

    resume_queue = CloudTasksQueueResumeOperator(
        location=LOCATION,
        queue_name=QUEUE_ID,
Ejemplo n.º 19
0
    CloudTasksTaskCreateOperator,
    CloudTasksTaskDeleteOperator,
    CloudTasksTaskGetOperator,
    CloudTasksTaskRunOperator,
    CloudTasksTasksListOperator,
)

GCP_CONN_ID = "google_cloud_default"
PROJECT_ID = "test-project"
LOCATION = "asia-east2"
FULL_LOCATION_PATH = "projects/test-project/locations/asia-east2"
QUEUE_ID = "test-queue"
FULL_QUEUE_PATH = "projects/test-project/locations/asia-east2/queues/test-queue"
TASK_NAME = "test-task"
FULL_TASK_PATH = "projects/test-project/locations/asia-east2/queues/test-queue/tasks/test-task"
TEST_QUEUE = Queue(name=FULL_QUEUE_PATH)
TEST_TASK = Task(app_engine_http_request={})


class TestCloudTasksQueueCreate(unittest.TestCase):
    @mock.patch("airflow.providers.google.cloud.operators.tasks.CloudTasksHook"
                )
    def test_create_queue(self, mock_hook):
        mock_hook.return_value.create_queue.return_value = TEST_QUEUE
        operator = CloudTasksQueueCreateOperator(location=LOCATION,
                                                 task_queue=TEST_QUEUE,
                                                 task_id="id")

        result = operator.execute(context=None)

        self.assertEqual({'name': FULL_QUEUE_PATH, 'state': 0}, result)