1 /*
2  * Copyright (C) 2024 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.tradefed.cache.remote;
18 
19 import static java.nio.charset.StandardCharsets.UTF_8;
20 import static org.junit.Assert.assertEquals;
21 import static org.junit.Assert.assertNull;
22 
23 import build.bazel.remote.execution.v2.ActionCacheGrpc.ActionCacheImplBase;
24 import build.bazel.remote.execution.v2.ActionResult;
25 import build.bazel.remote.execution.v2.Digest;
26 import build.bazel.remote.execution.v2.GetActionResultRequest;
27 import build.bazel.remote.execution.v2.UpdateActionResultRequest;
28 import com.android.tradefed.cache.DigestCalculator;
29 import com.android.tradefed.cache.ExecutableAction;
30 import com.android.tradefed.cache.ExecutableActionResult;
31 import com.android.tradefed.util.FileUtil;
32 import com.android.tradefed.util.StreamUtil;
33 import com.google.common.util.concurrent.Futures;
34 import com.google.common.util.concurrent.ListenableFuture;
35 import io.grpc.ManagedChannel;
36 import io.grpc.Server;
37 import io.grpc.Status;
38 import io.grpc.stub.StreamObserver;
39 import io.grpc.inprocess.InProcessChannelBuilder;
40 import io.grpc.inprocess.InProcessServerBuilder;
41 import io.grpc.util.MutableHandlerRegistry;
42 import java.io.ByteArrayInputStream;
43 import java.io.File;
44 import java.io.InputStream;
45 import java.io.IOException;
46 import java.io.OutputStream;
47 import java.time.Duration;
48 import java.util.Arrays;
49 import java.util.Collections;
50 import java.util.concurrent.TimeUnit;
51 import java.util.HashMap;
52 import java.util.Map;
53 import org.junit.After;
54 import org.junit.Before;
55 import org.junit.runner.RunWith;
56 import org.junit.runners.JUnit4;
57 import org.junit.Test;
58 
59 /** Tests for {@link RemoteCacheClient}. */
60 @RunWith(JUnit4.class)
61 public class RemoteCacheClientTest {
62     private static final String INSTANCE = "test instance";
63     private final String mFakeServerName = "fake server for " + getClass();
64     private final MutableHandlerRegistry mServiceRegistry = new MutableHandlerRegistry();
65     private ManagedChannel mChannel;
66     private Server mFakeServer;
67     private File mInput;
68     private File mWorkFolder;
69 
70     private static class FakeByteStreamDownloader extends ByteStreamDownloader {
71         private final Map<Digest, String> mData;
72 
FakeByteStreamDownloader(Map<Digest, String> data)73         public FakeByteStreamDownloader(Map<Digest, String> data) {
74             super(INSTANCE, null, null, Duration.ofSeconds(5));
75             mData = data;
76         }
77 
78         @Override
downloadBlob(Digest digest, OutputStream out)79         public ListenableFuture<Void> downloadBlob(Digest digest, OutputStream out) {
80             try {
81                 if (digest.getSizeBytes() == 0) {
82                     out.close();
83                     return Futures.immediateVoidFuture();
84                 }
85                 if (!mData.containsKey(digest)) {
86                     out.close();
87                     return Futures.immediateFailedFuture(new IOException("Blob not found!"));
88                 }
89                 InputStream data = new ByteArrayInputStream(mData.get(digest).getBytes(UTF_8));
90                 StreamUtil.copyStreams(data, out);
91                 out.close();
92                 data.close();
93             } catch (IOException e) {
94                 return Futures.immediateFailedFuture(e);
95             }
96             return Futures.immediateVoidFuture();
97         }
98     }
99 
100     @Before
setUp()101     public final void setUp() throws Exception {
102         mFakeServer =
103                 InProcessServerBuilder.forName(mFakeServerName)
104                         .fallbackHandlerRegistry(mServiceRegistry)
105                         .directExecutor()
106                         .build()
107                         .start();
108         mChannel = InProcessChannelBuilder.forName(mFakeServerName).directExecutor().build();
109         mInput = FileUtil.createTempDir("input-dir");
110         mWorkFolder = FileUtil.createTempDir("work-folder");
111     }
112 
113     @After
tearDown()114     public void tearDown() throws Exception {
115         mChannel.shutdown().awaitTermination(5, TimeUnit.SECONDS);
116         mFakeServer.shutdown().awaitTermination(5, TimeUnit.SECONDS);
117         FileUtil.recursiveDelete(mInput);
118         FileUtil.recursiveDelete(mWorkFolder);
119     }
120 
121     @Test
uploadCache_works()122     public void uploadCache_works() throws IOException, InterruptedException {
123         class SpyActionCacheImpl extends ActionCacheImplBase {
124             public ActionResult actionResult = null;
125             public Digest actionDigest = null;
126 
127             @Override
128             public void updateActionResult(
129                     UpdateActionResultRequest request,
130                     StreamObserver<ActionResult> responseObserver) {
131                 actionResult = request.getActionResult();
132                 actionDigest = request.getActionDigest();
133                 responseObserver.onNext(actionResult);
134                 responseObserver.onCompleted();
135             }
136         }
137         SpyActionCacheImpl actionCache = new SpyActionCacheImpl();
138         mServiceRegistry.addService(actionCache);
139         ExecutableAction action =
140                 ExecutableAction.create(
141                         mInput, Arrays.asList("test", "command"), new HashMap<>(), 100L);
142         int exitCode = 0;
143         File stdoutFile = FileUtil.createTempFile("stdout-", ".txt", mWorkFolder);
144         String stdout = "test stdout";
145         FileUtil.writeToFile(stdout, stdoutFile);
146         ExecutableActionResult result = ExecutableActionResult.create(exitCode, stdoutFile, null);
147         RemoteCacheClient client = newClient(null);
148         ActionResult expectedResult =
149                 ActionResult.newBuilder()
150                         .setExitCode(exitCode)
151                         .setStdoutDigest(DigestCalculator.compute(stdoutFile))
152                         .build();
153 
154         client.uploadCache(action, result);
155 
156         assertEquals(actionCache.actionResult, expectedResult);
157         assertEquals(actionCache.actionDigest, DigestCalculator.compute(action.action()));
158     }
159 
160     @Test
lookupCache_works()161     public void lookupCache_works() throws IOException, InterruptedException {
162         ExecutableAction notFoundAction =
163                 ExecutableAction.create(
164                         mInput, Arrays.asList("not", "found", "command"), new HashMap<>(), 100L);
165         ExecutableAction cachedAction =
166                 ExecutableAction.create(
167                         mInput, Arrays.asList("found", "command"), new HashMap<>(), 100L);
168         int exitCode = 0;
169         String stdout = "STDOUT";
170         Digest stdOutDigest = DigestCalculator.compute(stdout.getBytes());
171         mServiceRegistry.addService(
172                 new ActionCacheImplBase() {
173                     @Override
174                     public void getActionResult(
175                             GetActionResultRequest request,
176                             StreamObserver<ActionResult> responseObserver) {
177                         if (request.getActionDigest().equals(cachedAction.actionDigest())) {
178                             responseObserver.onNext(
179                                     ActionResult.newBuilder()
180                                             .setStdoutDigest(stdOutDigest)
181                                             .setExitCode(exitCode)
182                                             .build());
183                             responseObserver.onCompleted();
184                             return;
185                         }
186                         responseObserver.onError(Status.NOT_FOUND.asRuntimeException());
187                     }
188                 });
189         RemoteCacheClient client =
190                 newClient(
191                         new FakeByteStreamDownloader(
192                                 Collections.singletonMap(stdOutDigest, stdout)));
193 
194         ExecutableActionResult notFoundResult = client.lookupCache(notFoundAction);
195         ExecutableActionResult cachedResult = client.lookupCache(cachedAction);
196 
197         assertNull(notFoundResult);
198         assertEquals(0, cachedResult.exitCode());
199         assertEquals(stdout, FileUtil.readStringFromFile(cachedResult.stdOut()));
200         assertNull(cachedResult.stdErr());
201     }
202 
newClient(ByteStreamDownloader downloader)203     private RemoteCacheClient newClient(ByteStreamDownloader downloader) {
204         return new RemoteCacheClient(mWorkFolder, INSTANCE, mChannel, null, downloader);
205     }
206 }
207