scripts: add branch support to chrome_chromeos_lkgm

BUG=b:239106349
TEST=run_tests, manual runs

Change-Id: I5cd35049e32429591fd950b4f6b9c1f8f3add71c
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/3789787
Reviewed-by: Ben Pastene <bpastene@chromium.org>
Commit-Queue: Jack Neus <jackneus@google.com>
Tested-by: Jack Neus <jackneus@google.com>
diff --git a/scripts/chrome_chromeos_lkgm_unittest.py b/scripts/chrome_chromeos_lkgm_unittest.py
index 3a328e1..0774dfb 100644
--- a/scripts/chrome_chromeos_lkgm_unittest.py
+++ b/scripts/chrome_chromeos_lkgm_unittest.py
@@ -18,7 +18,8 @@
 
   def setUp(self):
     """Common set up method for all tests."""
-    self.committer = chrome_chromeos_lkgm.ChromeLKGMCommitter('1001.0.0')
+    self.committer = chrome_chromeos_lkgm.ChromeLKGMCommitter('1001.0.0',
+                                                              'main')
     self.old_lkgm = None
 
   @mock.patch('chromite.lib.gob_util.GetFileContentsOnHead')
@@ -51,7 +52,9 @@
   @mock.patch('chromite.lib.gob_util.GetFileContentsOnHead')
   def testVersionWithChromeBranch(self, mock_get_file):
     """Tests passing a version with a chrome branch strips the branch."""
-    self.committer = chrome_chromeos_lkgm.ChromeLKGMCommitter('1003.0.0-rc2')
+    branch = 'refs/branch-heads/5000'
+    self.committer = chrome_chromeos_lkgm.ChromeLKGMCommitter('1003.0.0-rc2',
+                                                              branch)
     mock_get_file.return_value = '1002.0.0'
 
     with mock.patch.object(self.committer._gerrit_helper, 'CreateChange') as cg:
@@ -62,6 +65,7 @@
           with mock.patch.object(self.committer._gerrit_helper, 'SetHashtags'):
             # Check the file was actually written out correctly.
             self.committer.UpdateLKGM()
+            cg.assert_called_once_with('chromium/src', branch, mock.ANY, False)
             ce.assert_called_once_with(123456, 'chromeos/CHROMEOS_LKGM',
                                        '1003.0.0')
             bc.assert_called_once_with(123456, labels={'Bot-Commit': 1},