Store BazelCapabilitiesClient in BazelNetwork

...and pass it to BazelCasClient

6 files changed, 34 insertions, 6 deletions

diff --git a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
index 67361036..347fa204 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
@@ -215,9 +215,11 @@ BazelCasClient::BazelCasClient(
     std::string const& server,
     Port port,
     gsl::not_null<Auth const*> const& auth,
-    gsl::not_null<RetryConfig const*> const& retry_config) noexcept
+    gsl::not_null<RetryConfig const*> const& retry_config,
+    gsl::not_null<BazelCapabilitiesClient const*> const& capabilities) noexcept
     : stream_{std::make_unique<ByteStreamClient>(server, port, auth)},
-      retry_config_{*retry_config} {
+      retry_config_{*retry_config},
+      capabilities_{*capabilities} {
     stub_ = bazel_re::ContentAddressableStorage::NewStub(
         CreateChannelWithCredentials(server, port, auth));
 }
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp
index 6e2f928b..0e719441 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp
@@ -36,6 +36,7 @@
 #include "src/buildtool/common/remote/retry_config.hpp"
 #include "src/buildtool/crypto/hash_function.hpp"
 #include "src/buildtool/execution_api/common/artifact_blob.hpp"
+#include "src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bytestream_client.hpp"
 #include "src/buildtool/logging/logger.hpp"
 
@@ -47,7 +48,9 @@ class BazelCasClient {
         std::string const& server,
         Port port,
         gsl::not_null<Auth const*> const& auth,
-        gsl::not_null<RetryConfig const*> const& retry_config) noexcept;
+        gsl::not_null<RetryConfig const*> const& retry_config,
+        gsl::not_null<BazelCapabilitiesClient const*> const&
+            capabilities) noexcept;
 
     /// \brief Find missing blobs
     /// \param[in] instance_name Name of the CAS instance
@@ -143,6 +146,7 @@ class BazelCasClient {
   private:
     std::unique_ptr<ByteStreamClient> stream_;
     RetryConfig const& retry_config_;
+    [[maybe_unused]] BazelCapabilitiesClient const& capabilities_;
     std::unique_ptr<bazel_re::ContentAddressableStorage::Stub> stub_;
     Logger logger_{"RemoteCasClient"};
 
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
index 823148f0..f7aa3dcc 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
@@ -31,7 +31,15 @@ BazelNetwork::BazelNetwork(
     ExecutionConfiguration const& exec_config,
     gsl::not_null<HashFunction const*> const& hash_function) noexcept
     : instance_name_{std::move(instance_name)},
-      cas_{std::make_unique<BazelCasClient>(host, port, auth, retry_config)},
+      capabilities_{std::make_unique<BazelCapabilitiesClient>(host,
+                                                              port,
+                                                              auth,
+                                                              retry_config)},
+      cas_{std::make_unique<BazelCasClient>(host,
+                                            port,
+                                            auth,
+                                            retry_config,
+                                            capabilities_.get())},
       ac_{std::make_unique<BazelAcClient>(host, port, auth, retry_config)},
       exec_{std::make_unique<BazelExecutionClient>(host,
                                                    port,
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
index 3b359700..1c65cbd7 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
@@ -31,6 +31,7 @@
 #include "src/buildtool/execution_api/bazel_msg/bazel_common.hpp"
 #include "src/buildtool/execution_api/common/artifact_blob.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_ac_client.hpp"
+#include "src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_execution_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp"
@@ -93,6 +94,7 @@ class BazelNetwork {
 
   private:
     std::string const instance_name_;
+    std::unique_ptr<BazelCapabilitiesClient> capabilities_;
     std::unique_ptr<BazelCasClient> cas_;
     std::unique_ptr<BazelAcClient> ac_;
     std::unique_ptr<BazelExecutionClient> exec_;
diff --git a/test/buildtool/execution_api/bazel/bazel_cas_client.test.cpp b/test/buildtool/execution_api/bazel/bazel_cas_client.test.cpp
index 2ec2a45f..cb2132b4 100644
--- a/test/buildtool/execution_api/bazel/bazel_cas_client.test.cpp
+++ b/test/buildtool/execution_api/bazel/bazel_cas_client.test.cpp
@@ -28,6 +28,7 @@
 #include "src/buildtool/common/remote/retry_config.hpp"
 #include "src/buildtool/crypto/hash_function.hpp"
 #include "src/buildtool/execution_api/common/artifact_blob.hpp"
+#include "src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.hpp"
 #include "src/buildtool/execution_api/remote/config.hpp"
 #include "src/buildtool/file_system/object_type.hpp"
 #include "test/utils/hermeticity/test_hash_function_type.hpp"
@@ -46,10 +47,15 @@ TEST_CASE("Bazel internals: CAS Client", "[execution_api]") {
     REQUIRE(remote_config);
     REQUIRE(remote_config->remote_address);
     RetryConfig retry_config{};  // default retry config
+    BazelCapabilitiesClient capabilities(remote_config->remote_address->host,
+                                         remote_config->remote_address->port,
+                                         &*auth_config,
+                                         &retry_config);
     BazelCasClient cas_client(remote_config->remote_address->host,
                               remote_config->remote_address->port,
                               &*auth_config,
-                              &retry_config);
+                              &retry_config,
+                              &capabilities);
 
     SECTION("Valid digest and blob") {
         // digest of "test"
diff --git a/test/utils/remote_execution/bazel_action_creator.hpp b/test/utils/remote_execution/bazel_action_creator.hpp
index 3ffcf827..49f29e9c 100644
--- a/test/utils/remote_execution/bazel_action_creator.hpp
+++ b/test/utils/remote_execution/bazel_action_creator.hpp
@@ -33,6 +33,7 @@
 #include "src/buildtool/common/remote/retry_config.hpp"
 #include "src/buildtool/crypto/hash_function.hpp"
 #include "src/buildtool/execution_api/common/artifact_blob.hpp"
+#include "src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp"
 #include "src/buildtool/execution_api/remote/config.hpp"
 #include "src/buildtool/file_system/object_type.hpp"
@@ -104,10 +105,15 @@
 
     RetryConfig retry_config{};  // default retry config
 
+    BazelCapabilitiesClient capabilities(remote_config->remote_address->host,
+                                         remote_config->remote_address->port,
+                                         &*auth_config,
+                                         &retry_config);
     BazelCasClient cas_client(remote_config->remote_address->host,
                               remote_config->remote_address->port,
                               &*auth_config,
-                              &retry_config);
+                              &retry_config,
+                              &capabilities);
 
     if (cas_client.BatchUpdateBlobs(instance_name, blobs) == blobs.size()) {
         return std::make_unique<bazel_re::Digest>(