service: Fix for get binhost locations

Update the service to call the binhost lookup service to get binhost
locations. This is a forward fix for:
https://chromium-review.googlesource.com/c/chromiumos/chromite/+/4832171

BUG=b:270985155
TEST=./run_tests service/binhost_unittest.py

Change-Id: Ia1307d80851016e815ef78d791a83350fc8dccbb
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/5121318
Commit-Queue: Nikhil Gumidelli <nikhilgm@google.com>
Tested-by: Nikhil Gumidelli <nikhilgm@google.com>
Reviewed-by: Greg Edelston <gredelston@google.com>
Reviewed-by: Cindy Lin <xcl@google.com>
diff --git a/service/binhost.py b/service/binhost.py
index e5314f8..3d1a763 100644
--- a/service/binhost.py
+++ b/service/binhost.py
@@ -46,15 +46,19 @@
 # Parameters for the Lookup Binhosts Service endpoint.
 _PROTOCOL = "https"
 _CHROMEOS_PREBUILTS_DOMAIN = "us-central1-chromeos-prebuilts.cloudfunctions.net"
-# TODO(b/272096443): Update to the production endpoint.
-_LOOKUP_BINHOSTS_ENDPOINT = "staging-public-lookup_service-binhosts"
-
-_LOOKUP_URI = "%s://%s/%s" % (
+_LOOKUP_BINHOSTS_ENDPOINT = "prod-lookup-service-binhosts"
+_BINHOST_LOOKUP_SERVICE_URI = "%s://%s/%s" % (
     _PROTOCOL,
     _CHROMEOS_PREBUILTS_DOMAIN,
     _LOOKUP_BINHOSTS_ENDPOINT,
 )
 
+# Timeout for the API call to the binhost lookup service.
+_BINHOST_LOOKUP_SERVICE_TIMEOUT = 60
+
+# Google storage bucket which contains binhosts.
+_BINHOSTS_GS_BUCKET_NAME = "chromeos-prebuilt"
+
 
 class Error(Exception):
     """Base error class for the module."""
@@ -72,6 +76,10 @@
     """When maximum number of uris to store is less than or equal to 0."""
 
 
+class BinhostsLookupServiceError(Error):
+    """When the binhost lookup service returns an error."""
+
+
 def _ValidateBinhostConf(path: Path, key: str) -> None:
     """Validates the binhost conf file defines only one environment variable.
 
@@ -661,23 +669,33 @@
     build_target: str,
     profile: str,
     private: bool,
+    get_corresponding_binhosts: bool,
+    generic_build_target: str,
+    generic_profile: str,
 ) -> List[Optional[str]]:
     """Call the binhost lookup service to get locations of BINHOSTs.
 
     Args:
         gs_bucket_name: Name of the google storage bucket which contains the
-            binhosts (e.g. "chromeos-prebuilts" in
-            gs://chromeos-prebuilts/binhosts/..).
+            binhosts (e.g. "chromeos-prebuilt" in
+            gs://chromeos-prebuilt/binhosts/..).
         snapshot_shas: List of snapshot shas of the binhosts.
         build_target: build target (also known as board) of the binhosts.
         profile: profile associated with the build target.
         private: True if the binhosts are private.
+        get_corresponding_binhosts: True if the corresponding internal/external
+            binhosts should also be returned.
+        generic_build_target: The base architecture board for the
+            build target (e.g. amd64-generic).
+        generic_profile: The profile of the base architecture board.
 
     Returns:
         A list of Google Storage URIs of binhosts, sorted by created
-        time (descending).
+        time (descending). Can be empty if no binhosts are found.
 
     Raises:
+        BinhostsLookupServiceError: When there's an error from the binhost
+            lookup service.
         google.protobuf.message.DecodeError: When the protobuf message from the
             API response cannot be parsed.
     """
@@ -688,7 +706,11 @@
         build_target=build_target,
         profile=profile,
         private=private,
+        get_corresponding_binhosts=get_corresponding_binhosts,
+        generic_build_target=generic_build_target,
+        generic_profile=generic_profile,
     )
+
     # Encode the request with base64 and convert to a string.
     lookup_binhosts_request_encoded = base64.urlsafe_b64encode(
         lookup_binhosts_request.SerializeToString()
@@ -696,14 +718,12 @@
 
     response = requests.request(
         "GET",
-        "%s://%s/%s?filter=%s"
+        "%s?filter=%s"
         % (
-            _PROTOCOL,
-            _CHROMEOS_PREBUILTS_DOMAIN,
-            _LOOKUP_BINHOSTS_ENDPOINT,
+            _BINHOST_LOOKUP_SERVICE_URI,
             lookup_binhosts_request_encoded,
         ),
-        timeout=70,
+        timeout=_BINHOST_LOOKUP_SERVICE_TIMEOUT,
     )
 
     if response.status_code == 200:
@@ -718,16 +738,14 @@
         logging.warning(
             "No suitable binhosts found in the binhost lookup service"
         )
+        return []
     else:
-        logging.error(
+        raise BinhostsLookupServiceError(
             "Error while fetching binhosts from the binhost lookup service, "
-            "status code: %i, body: %s",
-            response.status_code,
-            response.content,
+            "status code: %i, body: %s"
+            % (response.status_code, response.content)
         )
 
-    return []
-
 
 def _get_snapshot_shas() -> SnapshotShas:
     """Get snapshot SHAs for different checkout types.
@@ -746,7 +764,6 @@
         logging.error("Unable to determine a repo directory: %s", e)
         return snapshot_shas
 
-    # Determine the checkout types.
     manifest = repo.Manifest()
 
     # Get the snapshot SHAs for each checkout type.
@@ -795,14 +812,16 @@
 
 
 def lookup_binhosts(
-    gs_bucket_name: str, build_target: str, profile: str
+    build_target: str,
+    profile: str,
+    gs_bucket_name: str = _BINHOSTS_GS_BUCKET_NAME,
 ) -> List[Optional[str]]:
     """Get binhost locations from the binhost lookup service.
 
     Args:
         gs_bucket_name: Name of the google storage bucket which contains the
-            binhosts (e.g. "chromeos-prebuilts" in
-            gs://chromeos-prebuilts/binhosts/..).
+            binhosts (e.g. "chromeos-prebuilt" in
+            gs://chromeos-prebuilt/binhosts/..).
         build_target: build target (also known as board) of the binhosts.
         profile: profile associated with the build target.
 
@@ -810,34 +829,47 @@
         A list of Google Storage URIs of binhosts, sorted by created
         time (descending).
     """
-    snapshot_shas = _get_snapshot_shas()
-    binhost_gs_uris = []
+    snapshot_shas_combined = _get_snapshot_shas()
 
-    # Fetch internal binhosts.
-    if snapshot_shas.internal:
-        binhost_gs_uris.extend(
-            _fetch_binhosts(
-                gs_bucket_name,
-                snapshot_shas.internal,
-                build_target,
-                profile,
-                True,
-            )
-        )
+    site_params = config_lib.GetSiteParams()
+    # Get the repo.
+    try:
+        repo = repo_util.Repository.MustFind(constants.SOURCE_ROOT)
+    except repo_util.NotInRepoError as e:
+        logging.error("Unable to determine a repo directory: %s", e)
+        raise e
 
-    # Fetch external binhosts.
-    if snapshot_shas.external:
-        binhost_gs_uris.extend(
-            _fetch_binhosts(
-                gs_bucket_name,
-                snapshot_shas.external,
-                build_target,
-                profile,
-                False,
-            )
-        )
+    manifest = repo.Manifest()
 
-    # TODO(b/271207940): Handle use case for partners who need some external as
-    # well as internal binhosts.
+    if manifest.HasRemote(site_params.INTERNAL_REMOTE) and manifest.HasRemote(
+        site_params.EXTERNAL_REMOTE
+    ):
+        # Googlers
+        if snapshot_shas_combined.internal:
+            snapshot_shas = snapshot_shas_combined.internal
+            get_corresponding_binhosts = True
+            private = True
+        # Partners
+        else:
+            snapshot_shas = snapshot_shas_combined.external
+            get_corresponding_binhosts = True
+            private = False
+    # External Developers
+    else:
+        snapshot_shas = snapshot_shas_combined.external
+        get_corresponding_binhosts = False
+        private = False
+
+    # TODO(b/316157442): Map to the corresponding generic build target.
+    binhost_gs_uris = _fetch_binhosts(
+        gs_bucket_name,
+        snapshot_shas,
+        build_target,
+        profile,
+        private,
+        get_corresponding_binhosts,
+        build_target,
+        profile,
+    )
 
     return binhost_gs_uris
diff --git a/service/binhost_unittest.py b/service/binhost_unittest.py
index 7e38e26..eb7a5f1 100644
--- a/service/binhost_unittest.py
+++ b/service/binhost_unittest.py
@@ -36,6 +36,8 @@
 MOCK_BINHOST_ID = 1
 MOCK_BUILD_TARGET = "test_build_target"
 MOCK_DATE_STRING = "2023-07-25T08:09:14.842Z"
+MOCK_GENERIC_BUILD_TARGET = "generic_build_target"
+MOCK_GENERIC_PROFILE = "generic_profile"
 MOCK_GS_BUCKET_NAME = "test_bucket"
 MOCK_GS_URI = "gs://test"
 MOCK_ID_TOKEN = "test_token"
@@ -750,7 +752,7 @@
 class GetSnapshotShasTest(
     cros_test_lib.MockTempDirTestCase, cros_test_lib.LoggingTestCase
 ):
-    """The the _get_snapshot_shas function.
+    """Unittests for the _get_snapshot_shas function.
 
     The _get_snapshot_shas_from_git_log function is also tested here when it is
     called internally and the relevant functions are mocked.
@@ -867,6 +869,9 @@
         MOCK_BUILD_TARGET,
         MOCK_PROFILE,
         False,
+        True,
+        MOCK_GENERIC_BUILD_TARGET,
+        MOCK_GENERIC_PROFILE,
     )
 
     def setUp(self):
@@ -907,12 +912,8 @@
             400, "Unable to parse filter parameters from the request."
         )
 
-        with cros_test_lib.LoggingCapturer() as logs:
+        with self.assertRaises(binhost.BinhostsLookupServiceError):
             binhost._fetch_binhosts(*self.FETCH_BINHOSTS_MOCK_ARGS)
-            self.AssertLogsContain(
-                logs,
-                "Error while fetching binhosts from the binhost lookup service",
-            )
 
 
 class LookupBinhostsTest(cros_test_lib.MockTestCase):
@@ -933,17 +934,48 @@
         self.get_snapshot_shas = self.PatchObject(binhost, "_get_snapshot_shas")
         self.fetch_binhosts = self.PatchObject(binhost, "_fetch_binhosts")
 
-    def test_internal(self):
-        """Test for internal checkout having internal and external manifests."""
+    def test_internal_private_board(self):
+        """Test for internal checkout with a private board."""
         self.get_snapshot_shas.side_effect = [
             binhost.SnapshotShas(
                 self.EXTERNAL_SNAPSHOT_SHAS, self.INTERNAL_SNAPSHOT_SHAS
             )
         ]
-        self.fetch_binhosts.side_effect = (
-            self.INTERNAL_GS_URIS,
-            self.EXTERNAL_GS_URIS,
-        )
+        self.fetch_binhosts.side_effect = [
+            [
+                *self.INTERNAL_GS_URIS,
+                *self.EXTERNAL_GS_URIS,
+            ]
+        ]
+        assert binhost.lookup_binhosts(
+            MOCK_GS_BUCKET_NAME, MOCK_BUILD_TARGET, MOCK_PROFILE
+        ) == [*self.INTERNAL_GS_URIS, *self.EXTERNAL_GS_URIS]
+
+    def test_internal_public_board(self):
+        """Test for internal checkout with a public board."""
+        self.get_snapshot_shas.side_effect = [
+            binhost.SnapshotShas(
+                self.EXTERNAL_SNAPSHOT_SHAS, self.INTERNAL_SNAPSHOT_SHAS
+            )
+        ]
+        self.fetch_binhosts.side_effect = [self.EXTERNAL_GS_URIS]
+        assert binhost.lookup_binhosts(
+            MOCK_GS_BUCKET_NAME, MOCK_BUILD_TARGET, MOCK_PROFILE
+        ) == [*self.EXTERNAL_GS_URIS]
+
+    def test_external_private_board(self):
+        """Test for external checkout having internal and external manifests.
+
+        External checkouts which have both manifests but only have external
+        snapshot shas, e.g. partners.
+        """
+        self.get_snapshot_shas.side_effect = [
+            binhost.SnapshotShas(self.EXTERNAL_SNAPSHOT_SHAS, [])
+        ]
+        self.fetch_binhosts.side_effect = [
+            [*self.INTERNAL_GS_URIS, *self.EXTERNAL_GS_URIS]
+        ]
+
         assert binhost.lookup_binhosts(
             MOCK_GS_BUCKET_NAME, MOCK_BUILD_TARGET, MOCK_PROFILE
         ) == [*self.INTERNAL_GS_URIS, *self.EXTERNAL_GS_URIS]