/*
 * Copyright 2020 Google LLC
 *
 * 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.example.containeranalysis;

import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutures;
import com.google.api.gax.rpc.ApiCallContext;
import com.google.api.gax.rpc.UnaryCallable;
import com.google.cloud.devtools.containeranalysis.v1.ContainerAnalysisClient;
import io.grafeas.v1.GetOccurrenceRequest;
import io.grafeas.v1.GrafeasClient;
import io.grafeas.v1.NoteKind;
import io.grafeas.v1.Occurrence;
import io.grafeas.v1.Severity;
import io.grafeas.v1.VulnerabilityOccurrence;
import io.grafeas.v1.stub.GrafeasStub;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;

public class VulnerabilityFunctionTest {
  
  private final ContainerAnalysisClient containerAnalysisClient =
      Mockito.mock(ContainerAnalysisClient.class);
  private final GrafeasStub grafeasStub = Mockito.mock(GrafeasStub.class);
  private final GrafeasClient grafeasClient = Mockito.spy(GrafeasClient.create(grafeasStub));

  @Before
  public void setup() throws IOException {
    Mockito.when(containerAnalysisClient.getGrafeasClient()).thenReturn(grafeasClient);
  }

  @Test
  public void testAccept() throws IOException {
    AtomicInteger occurrenceCallCount = new AtomicInteger();

    Mockito.when(grafeasStub.getOccurrenceCallable())
        .thenReturn(
            new UnaryCallable<>() {
              @Override
              public ApiFuture<Occurrence> futureCall(
                  GetOccurrenceRequest request, ApiCallContext context) {
                occurrenceCallCount.incrementAndGet();

                return ApiFutures.immediateFuture(
                    Occurrence.newBuilder()
                        .setKind(NoteKind.VULNERABILITY)
                        .setResourceUri("gcr.io/test-project/some-image")
                        .setVulnerability(
                            VulnerabilityOccurrence.newBuilder()
                                .setSeverity(Severity.CRITICAL)
                                .setShortDescription("CVE-CRITICAL")
                                .build())
                        .build());
              }
            });

    VulnerabilityFunction function = new VulnerabilityFunction(containerAnalysisClient);

    String notificationPayload =
        "{\"name\":\"projects/test-project/occurrences/some-uuid\",\"kind\":\"VULNERABILITY\","
            + "\"notificationTime\":\"2020-09-04T00:38:25.575543Z\"}";
    PubSubMessage msg = new PubSubMessage();
    msg.setData(
        Base64.getEncoder().encodeToString(notificationPayload.getBytes(StandardCharsets.UTF_8)));
    function.accept(msg, null);

    Assert.assertEquals(1, occurrenceCallCount.get());
  }
}