googleapis · parthea · Oct 3, 2022 · Sep 30, 2022 · Oct 3, 2022 · Oct 3, 2022
diff --git a/samples/snippets/create/create_with_container_no_mounting.py b/samples/snippets/create/create_with_container_no_mounting.py
@@ -54,6 +54,7 @@ def create_container_job(project_id: str, region: str, job_name: str) -> batch_v
     task.max_run_duration = "3600s"
 
     # Tasks are grouped inside a job using TaskGroups.
+    # Currently, it's possible to have only one task group.
     group = batch_v1.TaskGroup()
     group.task_count = 4
     group.task_spec = task

diff --git a/samples/snippets/create/create_with_mounted_bucket.py b/samples/snippets/create/create_with_mounted_bucket.py
@@ -58,6 +58,7 @@ def create_script_job_with_bucket(project_id: str, region: str, job_name: str, b
     task.max_run_duration = "3600s"
 
     # Tasks are grouped inside a job using TaskGroups.
+    # Currently, it's possible to have only one task group.
     group = batch_v1.TaskGroup()
     group.task_count = 4
     group.task_spec = task

diff --git a/samples/snippets/create/create_with_script_no_mounting.py b/samples/snippets/create/create_with_script_no_mounting.py
@@ -54,6 +54,7 @@ def create_script_job(project_id: str, region: str, job_name: str) -> batch_v1.J
     task.max_run_duration = "3600s"
 
     # Tasks are grouped inside a job using TaskGroups.
+    # Currently, it's possible to have only one task group.
     group = batch_v1.TaskGroup()
     group.task_count = 4
     group.task_spec = task

diff --git a/samples/snippets/create/create_with_template.py b/samples/snippets/create/create_with_template.py
@@ -57,6 +57,7 @@ def create_script_job_with_template(project_id: str, region: str, job_name: str,
     task.max_run_duration = "3600s"
 
     # Tasks are grouped inside a job using TaskGroups.
+    # Currently, it's possible to have only one task group.
     group = batch_v1.TaskGroup()
     group.task_count = 4
     group.task_spec = task

diff --git a/samples/snippets/get/get_task.py b/samples/snippets/get/get_task.py
@@ -0,0 +1,38 @@
+#  Copyright 2022 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+# [START batch_get_task]
+
+from google.cloud import batch_v1
+
+
+def get_task(project_id: str, region: str, job_name: str, group_name: str, task_number: int) -> batch_v1.Task:
+    """
+    Retrieve information about a Task.
+
+    Args:
+        project_id: project ID or project number of the Cloud project you want to use.
+        region: name of the region hosts the job.
+        job_name: the name of the job you want to retrieve information about.
+        group_name: the name of the group that owns the task you want to check. Usually it's `group0`.
+        task_number: number of the task you want to look up.
+
+    Returns:
+        A Task object representing the specified task.
+    """
+    client = batch_v1.BatchServiceClient()
+
+    return client.get_task(name=f"projects/{project_id}/locations/{region}/jobs/{job_name}"
+                                f"/taskGroups/{group_name}/tasks/{task_number}")
+# [END batch_get_task]
diff --git a/samples/snippets/list/list_jobs.py b/samples/snippets/list/list_jobs.py
@@ -11,9 +11,10 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-from typing import Iterable
 
 # [START batch_list_jobs]
+from typing import Iterable
+
 from google.cloud import batch_v1
 
 

diff --git a/samples/snippets/list/list_tasks.py b/samples/snippets/list/list_tasks.py
@@ -0,0 +1,37 @@
+#  Copyright 2022 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+# [START batch_list_tasks]
+from typing import Iterable
+
+from google.cloud import batch_v1
+
+
+def list_tasks(project_id: str, region: str, job_name: str, group_name: str) -> Iterable[batch_v1.Task]:
+    """
+    Get a list of all jobs defined in given region.
+
+    Args:
+        project_id: project ID or project number of the Cloud project you want to use.
+        region: name of the region hosting the jobs.
+        job_name: name of the job which tasks you want to list.
+        group_name: name of the group of tasks. Usually it's `group0`.
+
+    Returns:
+        An iterable collection of Task objects.
+    """
+    client = batch_v1.BatchServiceClient()
+
+    return client.list_tasks(parent=f"projects/{project_id}/locations/{region}/jobs/{job_name}/taskGroups/{group_name}")
+# [END batch_list_tasks]
diff --git a/samples/snippets/tests/test_basics.py b/samples/snippets/tests/test_basics.py
@@ -24,7 +24,9 @@
 
 from ..delete.delete_job import delete_job
 from ..get.get_job import get_job
+from ..get.get_task import get_task
 from ..list.list_jobs import list_jobs
+from ..list.list_tasks import list_tasks
 
 PROJECT = google.auth.default()[1]
 REGION = 'europe-north1'
@@ -36,6 +38,7 @@
     batch_v1.JobStatus.State.QUEUED,
     batch_v1.JobStatus.State.RUNNING,
     batch_v1.JobStatus.State.SCHEDULED,
+    batch_v1.JobStatus.State.DELETION_IN_PROGRESS
 }
 
 
@@ -53,8 +56,7 @@ def _test_body(test_job: batch_v1.Job, additional_test: Callable = None):
             test_job = get_job(PROJECT, REGION, test_job.name.rsplit('/', maxsplit=1)[1])
             time.sleep(5)
 
-        assert test_job.status.state in (batch_v1.JobStatus.State.SUCCEEDED,
-                                         batch_v1.JobStatus.State.DELETION_IN_PROGRESS)
+        assert test_job.status.state == batch_v1.JobStatus.State.SUCCEEDED
 
         for job in list_jobs(PROJECT, REGION):
             if test_job.uid == job.uid:
@@ -72,9 +74,17 @@ def _test_body(test_job: batch_v1.Job, additional_test: Callable = None):
             pytest.fail("The test job should be deleted at this point!")
 
 
+def _check_tasks(job_name):
+    tasks = list_tasks(PROJECT, REGION, job_name, 'group0')
+    assert len(list(tasks)) == 4
+    for i in range(4):
+        assert get_task(PROJECT, REGION, job_name, 'group0', i) is not None
+    print('Tasks tested')
+
+
 def test_script_job(job_name):
     job = create_script_job(PROJECT, REGION, job_name)
-    _test_body(job)
+    _test_body(job, additional_test=lambda: _check_tasks(job_name))
 
 
 def test_container_job(job_name):