/*
 * Copyright 2019-present Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may
 * not use this file except in compliance with the License. You may obtain
 * a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations
 * under the License.
 */
package com.facebook.buck.remoteexecution.grpc;

import build.bazel.remote.execution.v2.ExecuteRequest;
import build.bazel.remote.execution.v2.ExecutionGrpc;
import build.bazel.remote.execution.v2.ExecutionGrpc.ExecutionBlockingStub;
import build.bazel.remote.execution.v2.ExecutionGrpc.ExecutionImplBase;
import com.facebook.buck.remoteexecution.grpc.retry.RetryClientInterceptor;
import com.facebook.buck.remoteexecution.grpc.retry.RetryPolicy;
import com.google.longrunning.Operation;
import io.grpc.ManagedChannel;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import io.grpc.inprocess.InProcessChannelBuilder;
import io.grpc.inprocess.InProcessServerBuilder;
import io.grpc.stub.StreamObserver;
import java.io.IOException;
import org.junit.Assert;
import org.junit.Test;

public class GrpcRetryInterceptorTest {
  private static class ExecutionImpl extends ExecutionImplBase {
    public int calls = 0;
    private final Status status;

    public ExecutionImpl(Status status) {
      this.status = status;
    }

    @Override
    public void execute(ExecuteRequest request, StreamObserver<Operation> responseObserver) {
      calls++;
      responseObserver.onError(status.asRuntimeException());
    }
  }

  private static class CallCounter implements Runnable {
    public int calls = 0;

    @Override
    public void run() {
      calls++;
    }
  }

  @Test
  public void testRetryOnUnavailable() throws IOException {
    String uniqueName = InProcessServerBuilder.generateName();
    ExecutionImpl service = new ExecutionImpl(Status.UNAVAILABLE);
    InProcessServerBuilder.forName(uniqueName).addService(service).build().start();
    CallCounter beforeRetry = new CallCounter();
    ManagedChannel channel =
        InProcessChannelBuilder.forName(uniqueName)
            .intercept(
                new RetryClientInterceptor(
                    RetryPolicy.builder().setMaxRetries(2).setBeforeRetry(beforeRetry).build()))
            .build();
    ExecutionBlockingStub stub = ExecutionGrpc.newBlockingStub(channel);
    try {
      stub.execute(ExecuteRequest.newBuilder().build()).forEachRemaining(resp -> {});
      Assert.fail("Final retry should cause an exception");
    } catch (StatusRuntimeException ex) {
      Assert.assertEquals(Status.Code.UNAVAILABLE, ex.getStatus().getCode());
    }

    Assert.assertEquals(3, service.calls);
    Assert.assertEquals(2, beforeRetry.calls);
  }

  @Test
  public void testNoRetryOnOk() throws IOException {
    String uniqueName = InProcessServerBuilder.generateName();
    ExecutionImpl service = new ExecutionImpl(Status.OK);
    InProcessServerBuilder.forName(uniqueName).addService(service).build().start();
    CallCounter beforeRetry = new CallCounter();
    ManagedChannel channel =
        InProcessChannelBuilder.forName(uniqueName)
            .intercept(
                new RetryClientInterceptor(
                    RetryPolicy.builder().setMaxRetries(2).setBeforeRetry(beforeRetry).build()))
            .build();
    ExecutionBlockingStub stub = ExecutionGrpc.newBlockingStub(channel);
    stub.execute(ExecuteRequest.newBuilder().build()).forEachRemaining(resp -> {});

    Assert.assertEquals(1, service.calls);
    Assert.assertEquals(0, beforeRetry.calls);
  }
}

@@ -1496,16 +1496,57 @@
 plBase %7B
+%0A%0A    private final int partialResponses;
 %0A    pub
@@ -1636,16 +1636,38 @@
 s status
+, int partialResponses
 ) %7B%0A    
@@ -1690,16 +1690,64 @@
 status;%0A
+      this.partialResponses = partialResponses;%0A
     %7D%0A%0A 
@@ -1864,24 +1864,148 @@
    calls++;%0A
+      for (int i = 0; i %3C partialResponses; i++) %7B%0A        responseObserver.onNext(Operation.newBuilder().build());%0A      %7D%0A
       respon
@@ -2407,16 +2407,19 @@
 VAILABLE
+, 0
 );%0A    I
@@ -3444,16 +3444,19 @@
 tatus.OK
+, 0
 );%0A    I
@@ -4102,11 +4102,1202 @@
 ls);%0A  %7D
+%0A%0A  @Test%0A  public void testRetryOnPartialResponse() throws IOException %7B%0A    String uniqueName = InProcessServerBuilder.generateName();%0A    ExecutionImpl service = new ExecutionImpl(Status.UNAVAILABLE, 1);%0A    InProcessServerBuilder.forName(uniqueName).addService(service).build().start();%0A    CallCounter beforeRetry = new CallCounter();%0A    ManagedChannel channel =%0A        InProcessChannelBuilder.forName(uniqueName)%0A            .intercept(%0A                new RetryClientInterceptor(%0A                    RetryPolicy.builder()%0A                        .setMaxRetries(2)%0A                        .setBeforeRetry(beforeRetry)%0A                        .setRestartAllStreamingCalls(true)%0A                        .build()))%0A            .build();%0A    ExecutionBlockingStub stub = ExecutionGrpc.newBlockingStub(channel);%0A    try %7B%0A      stub.execute(ExecuteRequest.newBuilder().build()).forEachRemaining(resp -%3E %7B%7D);%0A      Assert.fail(%22Final retry should cause an exception%22);%0A    %7D catch (StatusRuntimeException ex) %7B%0A      Assert.assertEquals(Status.Code.UNAVAILABLE, ex.getStatus().getCode());%0A    %7D%0A%0A    Assert.assertEquals(3, service.calls);%0A    Assert.assertEquals(2, beforeRetry.calls);%0A  %7D
 %0A%7D%0A

verify that partial responses are retried on setRestartAllStreamingCalls