buildbucket_v2: Add schedule, update, and cancel functionality

Add functionality to schedule, update, and cancel builds via the
Buildbucket V2 API.

This is the first step in removing V1 functionality, replaced with V2
for ongoing support.

BUG=b:179735986
TEST=`cros tryjob amd64-generic-full-tryjob`

Change-Id: I296e958f7f7740257778fdefe5f582eaf7715976
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2683055
Reviewed-by: LaMont Jones <lamontjones@chromium.org>
Commit-Queue: Mike Nichols <mikenichols@chromium.org>
Tested-by: Mike Nichols <mikenichols@chromium.org>
diff --git a/lib/buildbucket_v2.py b/lib/buildbucket_v2.py
index 8af548a..2b91656 100644
--- a/lib/buildbucket_v2.py
+++ b/lib/buildbucket_v2.py
@@ -245,25 +245,112 @@
   @retry_util.WithRetry(max_retry=5, sleep=20.0, exception=socket.error)
   @retry_util.WithRetry(max_retry=5, sleep=20.0,
                         exception=httplib.ResponseNotReady)
-  def GetBuild(self, buildbucket_id, properties=None):
-    """GetBuild call of a specific build with buildbucket_id.
+  def CancelBuild(self, buildbucket_id, summary_markdown, properties=None):
+    """CancelBuild call of a specific build with buildbucket_id.
 
     Args:
       buildbucket_id: id of the build in buildbucket.
-      properties: specific build.output.properties to query.
+      summary_markdown: Human-readable summary of the build in Markdown format.
+      properties: fields to include in the response.
 
     Returns:
       The corresponding Build proto. See here:
       https://chromium.googlesource.com/infra/luci/luci-go/+/master/buildbucket/proto/build.proto
     """
-    if properties:
-      field_mask = field_mask_pb2.FieldMask(paths=[properties])
-      get_build_request = rpc_pb2.GetBuildRequest(id=buildbucket_id,
-                                                  fields=field_mask)
-    else:
-      get_build_request = rpc_pb2.GetBuildRequest(id=buildbucket_id)
+    cancel_build_request = rpc_pb2.CancelBuildRequest(
+         id=buildbucket_id,
+         summary_markdown=summary_markdown,
+         fields=(field_mask_pb2.FieldMask(paths=[properties])
+                 if properties else None)
+    )
+    return self.client.CancelBuild(cancel_build_request)
+
+  # TODO(crbug/1006818): Need to handle ResponseNotReady given by luci prpc.
+  @retry_util.WithRetry(max_retry=5, sleep=20.0, exception=SSLError)
+  @retry_util.WithRetry(max_retry=5, sleep=20.0, exception=socket.error)
+  @retry_util.WithRetry(max_retry=5, sleep=20.0,
+                        exception=httplib.ResponseNotReady)
+  def GetBuild(self, buildbucket_id, properties=None):
+    """GetBuild call of a specific build with buildbucket_id.
+
+    Args:
+      buildbucket_id: id of the build in buildbucket.
+      properties: fields to include in the response.
+
+    Returns:
+      The corresponding Build proto. See here:
+      https://chromium.googlesource.com/infra/luci/luci-go/+/master/buildbucket/proto/build.proto
+    """
+    get_build_request = rpc_pb2.GetBuildRequest(
+        id=buildbucket_id,
+        fields=(field_mask_pb2.FieldMask(paths=[properties])
+                if properties else None)
+    )
     return self.client.GetBuild(get_build_request)
 
+# TODO(crbug/1006818): Need to handle ResponseNotReady given by luci prpc.
+  @retry_util.WithRetry(max_retry=5, sleep=20.0, exception=SSLError)
+  @retry_util.WithRetry(max_retry=5, sleep=20.0, exception=socket.error)
+  @retry_util.WithRetry(max_retry=5, sleep=20.0,
+                        exception=httplib.ResponseNotReady)
+  def ScheduleBuild(self, request_id, template_build_id=None,
+                    builder=None, experiments=None,
+                    properties=None, gerrit_changes=None,
+                    tags=None, fields=None, critical=True):
+    """GetBuild call of a specific build with buildbucket_id.
+
+    Args:
+      request_id: unique string used to prevent duplicates.
+      template_build_id: ID of a build to retry.
+      builder: Tuple (builder.project, builder.bucket) defines build ACL
+      experiments: map of string, bool of experiments to set.
+      properties: properties key in parameters_json
+      gerrit_changes: Repeated GerritChange message type.
+      tags: repeated StringPair of Build.tags to associate with build.
+      fields: fields to include in the response.
+      critical: bool for build.critical.
+
+    Returns:
+      The corresponding Build proto. See here:
+      https://chromium.googlesource.com/infra/luci/luci-go/+/master/buildbucket/proto/build.proto
+    """
+    schedule_build_request = rpc_pb2.ScheduleBuildRequest(
+        request_id=request_id,
+        template_build_id=template_build_id if template_build_id else None,
+        builder=builder,
+        experiments=experiments if experiments else None,
+        properties=properties if properties else None,
+        gerrit_changes=gerrit_changes if gerrit_changes else None,
+        tags=tags if tags else None,
+        fields=field_mask_pb2.FieldMask(paths=[fields]) if fields else None,
+        critical=common_pb2.YES if critical else common_pb2.NO)
+    return self.client.ScheduleBuild(schedule_build_request)
+
+# TODO(crbug/1006818): Need to handle ResponseNotReady given by luci prpc.
+  @retry_util.WithRetry(max_retry=5, sleep=20.0, exception=SSLError)
+  @retry_util.WithRetry(max_retry=5, sleep=20.0, exception=socket.error)
+  @retry_util.WithRetry(max_retry=5, sleep=20.0,
+                        exception=httplib.ResponseNotReady)
+  def UpdateBuild(self, build, update_properties, properties=None):
+    """GetBuild call of a specific build with buildbucket_id.
+
+    Args:
+      build: Buildbucket build to update.
+      update_properties: fields to update.
+      properties: fields to include in the response.
+
+    Returns:
+      The corresponding Build proto. See here:
+      https://chromium.googlesource.com/infra/luci/luci-go/+/master/buildbucket/proto/build.proto
+    """
+    update_build_request = rpc_pb2.UpdateBuildRequest(
+        build=build,
+        update_mask=field_mask_pb2.FieldMask(paths=[update_properties]),
+        fields=(field_mask_pb2.FieldMask(paths=[properties])
+                if properties else None)
+    )
+    return self.client.UpdateBuild(update_build_request)
+
   def GetKilledChildBuilds(self, buildbucket_id):
     """Get IDs of all the builds killed by self-destructed master build.
 
diff --git a/lib/buildbucket_v2_unittest.py b/lib/buildbucket_v2_unittest.py
index 38d8665..4d965b7 100644
--- a/lib/buildbucket_v2_unittest.py
+++ b/lib/buildbucket_v2_unittest.py
@@ -36,6 +36,34 @@
     ret = buildbucket_v2.BuildbucketV2(test_env=True)
     self.assertIsInstance(ret.client, Client)
 
+  def testCancelBuildWithProperties(self):
+    fake_field_mask = field_mask_pb2.FieldMask(paths=['properties'])
+    fake_cancel_build_request = object()
+    bbv2 = buildbucket_v2.BuildbucketV2()
+    client = bbv2.client
+    self.cancel_build_request_fn = self.PatchObject(
+        rpc_pb2, 'CancelBuildRequest', return_value=fake_cancel_build_request)
+    self.cancel_build_function = self.PatchObject(client, 'CancelBuild')
+    bbv2.CancelBuild('some-id', 'summary_markdown', 'properties')
+    self.cancel_build_request_fn.assert_called_with(id='some-id',
+        summary_markdown='summary_markdown',
+        fields=fake_field_mask)
+    self.cancel_build_function.assert_called_with(fake_cancel_build_request)
+
+  def testCancelBuildWithoutProperties(self):
+    fake_cancel_build_request = object()
+    bbv2 = buildbucket_v2.BuildbucketV2()
+    client = bbv2.client
+    self.cancel_build_request_fn = self.PatchObject(
+        rpc_pb2, 'CancelBuildRequest', return_value=fake_cancel_build_request)
+    self.cancel_build_function = self.PatchObject(client, 'CancelBuild')
+    bbv2.CancelBuild('some-id', 'summary_markdown')
+    self.cancel_build_request_fn.assert_called_with(
+        id='some-id',
+        summary_markdown='summary_markdown',
+        fields=None)
+    self.cancel_build_function.assert_called_with(fake_cancel_build_request)
+
   def testGetBuildWithProperties(self):
     fake_field_mask = field_mask_pb2.FieldMask(paths=['properties'])
     fake_get_build_request = object()
@@ -57,9 +85,72 @@
         rpc_pb2, 'GetBuildRequest', return_value=fake_get_build_request)
     self.get_build_function = self.PatchObject(client, 'GetBuild')
     bbv2.GetBuild('some-id')
-    self.get_build_request_fn.assert_called_with(id='some-id')
+    self.get_build_request_fn.assert_called_with(id='some-id', fields=None)
     self.get_build_function.assert_called_with(fake_get_build_request)
 
+  def testScheduleBuild(self):
+    fake_builder = build_pb2.BuilderID(project='chromeos',
+                                      bucket='general',
+                                      builder='test-builder')
+    fake_field_mask = field_mask_pb2.FieldMask(paths=['properties'])
+    fake_tag = common_pb2.StringPair(key='foo',
+                                     value='bar')
+    fake_schedule_build_request = object()
+    bbv2 = buildbucket_v2.BuildbucketV2()
+    client = bbv2.client
+    self.schedule_build_request_fn = self.PatchObject(
+        rpc_pb2, 'ScheduleBuildRequest',
+        return_value=fake_schedule_build_request)
+    self.schedule_build_function = self.PatchObject(client, 'ScheduleBuild')
+    bbv2.ScheduleBuild(request_id='1234', builder=fake_builder,
+                       tags=fake_tag, fields='properties', critical=True)
+    self.schedule_build_request_fn.assert_called_with(
+        request_id='1234',
+        template_build_id=None,
+        builder=fake_builder,
+        experiments=None,
+        properties=None,
+        gerrit_changes=None,
+        tags=fake_tag,
+        fields=fake_field_mask,
+        critical=common_pb2.YES)
+    self.schedule_build_function.assert_called_with(fake_schedule_build_request)
+
+  def testUpdateBuildWithProperties(self):
+    fake_update_mask = field_mask_pb2.FieldMask(paths=['number'])
+    fake_prop_mask = field_mask_pb2.FieldMask(paths=['properties'])
+    fake_build = build_pb2.Build(id=2341,
+                                 number=1234)
+    fake_update_build_request = object()
+    bbv2 = buildbucket_v2.BuildbucketV2()
+    client = bbv2.client
+    self.update_build_request_fn = self.PatchObject(
+        rpc_pb2, 'UpdateBuildRequest', return_value=fake_update_build_request)
+    self.update_build_function = self.PatchObject(client, 'UpdateBuild')
+    bbv2.UpdateBuild(fake_build, 'number', 'properties')
+    self.update_build_request_fn.assert_called_with(
+        build=fake_build,
+        update_mask=fake_update_mask,
+        fields=fake_prop_mask)
+    self.update_build_function.assert_called_with(fake_update_build_request)
+
+  def testUpdateBuildWithoutProperties(self):
+    fake_update_mask = field_mask_pb2.FieldMask(paths=['number'])
+    fake_build = build_pb2.Build(id=2341,
+                                 number=1234)
+    fake_update_build_request = object()
+    bbv2 = buildbucket_v2.BuildbucketV2()
+    client = bbv2.client
+    self.update_build_request_fn = self.PatchObject(
+        rpc_pb2, 'UpdateBuildRequest', return_value=fake_update_build_request)
+    self.update_build_function = self.PatchObject(client, 'UpdateBuild')
+    bbv2.UpdateBuild(fake_build, 'number')
+    self.update_build_request_fn.assert_called_with(
+        build=fake_build,
+        update_mask=fake_update_mask,
+        fields=None)
+    self.update_build_function.assert_called_with(fake_update_build_request)
+
   def testGetBuildStages(self):
     """Test the GetBuildStages functionality."""
     bbv2 = buildbucket_v2.BuildbucketV2()