From 9ead1e019cad2c5ecf3d052cfe07cfd3934ffe06 Mon Sep 17 00:00:00 2001
From: Miles Johnson <mileswjohnson@gmail.com>
Date: Thu, 23 Jan 2025 18:24:25 -0800
Subject: [PATCH] Fix tests.

---
 crates/remote/src/fs_digest.rs               |  3 ++-
 crates/remote/src/grpc_remote_client.rs      |  2 --
 crates/remote/src/remote_service.rs          |  6 ++++--
 crates/task-runner/tests/task_runner_test.rs | 21 +++-----------------
 4 files changed, 9 insertions(+), 23 deletions(-)

diff --git a/crates/remote/src/fs_digest.rs b/crates/remote/src/fs_digest.rs
index e2eb5268ea9..8b07edd3236 100644
--- a/crates/remote/src/fs_digest.rs
+++ b/crates/remote/src/fs_digest.rs
@@ -3,7 +3,7 @@
 use bazel_remote_apis::build::bazel::remote::execution::v2::{
     Digest, NodeProperties, OutputDirectory, OutputFile, OutputSymlink,
 };
-use bazel_remote_apis::google::protobuf::{Timestamp, UInt32Value};
+use bazel_remote_apis::google::protobuf::Timestamp;
 use chrono::NaiveDateTime;
 use moon_common::path::{PathExt, WorkspaceRelativePathBuf};
 use sha2::{Digest as Sha256Digest, Sha256};
@@ -76,6 +76,7 @@ pub fn compute_node_properties(metadata: &Metadata) -> NodeProperties {
 
     #[cfg(unix)]
     {
+        use bazel_remote_apis::google::protobuf::UInt32Value;
         use std::os::unix::fs::PermissionsExt;
 
         props.unix_mode = Some(UInt32Value {
diff --git a/crates/remote/src/grpc_remote_client.rs b/crates/remote/src/grpc_remote_client.rs
index bf241e5f678..0034003474f 100644
--- a/crates/remote/src/grpc_remote_client.rs
+++ b/crates/remote/src/grpc_remote_client.rs
@@ -22,14 +22,12 @@ use tonic::{
 use tracing::{trace, warn};
 
 fn map_transport_error(error: tonic::transport::Error) -> RemoteError {
-    dbg!(&error);
     RemoteError::GrpcConnectFailed {
         error: Box::new(error),
     }
 }
 
 fn map_status_error(error: tonic::Status) -> RemoteError {
-    dbg!(&error);
     RemoteError::GrpcCallFailed {
         error: Box::new(error),
     }
diff --git a/crates/remote/src/remote_service.rs b/crates/remote/src/remote_service.rs
index 4a3ebeb2dd9..4cc0b571755 100644
--- a/crates/remote/src/remote_service.rs
+++ b/crates/remote/src/remote_service.rs
@@ -161,8 +161,10 @@ impl RemoteService {
                     Some(cap.max_batch_total_size_bytes)
                 }
             })
-            // grpc limit: 4mb - buffer
-            .unwrap_or(4194304 - (1024 * 10))
+            // grpc limit: 4mb
+            .unwrap_or(4194304)
+            // Minus a small buffer
+            - 1024
     }
 
     #[instrument(skip(self, state))]
diff --git a/crates/task-runner/tests/task_runner_test.rs b/crates/task-runner/tests/task_runner_test.rs
index 2109bf539a3..8291994dce4 100644
--- a/crates/task-runner/tests/task_runner_test.rs
+++ b/crates/task-runner/tests/task_runner_test.rs
@@ -726,11 +726,7 @@ mod task_runner {
             runner.execute(&context, &node).await.unwrap();
 
             assert_eq!(
-                context
-                    .target_states
-                    .get(&runner.task.target)
-                    .unwrap()
-                    .get(),
+                runner.target_state.as_ref().unwrap(),
                 &TargetState::Passed("hash123".into())
             );
         }
@@ -747,11 +743,7 @@ mod task_runner {
             runner.execute(&context, &node).await.unwrap();
 
             assert_eq!(
-                context
-                    .target_states
-                    .get(&runner.task.target)
-                    .unwrap()
-                    .get(),
+                runner.target_state.as_ref().unwrap(),
                 &TargetState::Passthrough
             );
         }
@@ -770,14 +762,7 @@ mod task_runner {
             // Swallow panic so we can check operations
             let _ = runner.execute(&context, &node).await;
 
-            assert_eq!(
-                context
-                    .target_states
-                    .get(&runner.task.target)
-                    .unwrap()
-                    .get(),
-                &TargetState::Failed
-            );
+            assert_eq!(runner.target_state.as_ref().unwrap(), &TargetState::Failed);
         }
 
         #[tokio::test]