使用 JUnit 模拟 Google.Cloud.AIPlatform.V1 中的应用程序默认凭证

java

2024-03-04 11:24:49

使用 JUnit 模拟 google.cloud.aiplatform.v1 中的应用程序默认凭证

前言

在单元测试中模拟应用程序默认凭证 (ADC) 至关重要，以确保测试反映实际应用程序行为。本文将指导你逐步模拟 google.cloud.aiplatform.v1 中的 ADC，提升测试效率和可靠性。

背景

ADC 是一种用于向 Google Cloud API 认证应用程序的凭证。如果没有设置 "GOOGLE_APPLICATION_CREDENTIALS" 环境变量，应用程序在访问受保护的 API 时将抛出错误。

模拟 ADC

1. 引入依赖项

import com.google.auth.oauth2.GoogleCredentials;

2. 创建模拟凭证

public class MockGoogleCredentials extends GoogleCredentials {
    @Override public GoogleCredentials createScopedRequired() { return this; }
    @Override public AccessToken refreshAccessToken() { return new AccessToken("", null); }
}

3. 注入模拟凭证

@Before
public void setUp() {
    GoogleCredentials mockCredentials = new MockGoogleCredentials();
    when(GoogleCredentials.getApplicationDefault()).thenReturn(mockCredentials);
}

4. 编写测试

@Test
public void testGetVertexAIPredictor() {
    PredictorV2 predictorV2 = vertexAIPredictorProvider.getPredictor("vertexai");
    assertNotNull(predictorV2);
}

完整代码示例

import com.google.api.gax.longrunning.OperationFuture;
import com.google.cloud.aiplatform.v1.EndpointName;
import com.google.cloud.aiplatform.v1.PredictResponse;
import com.google.cloud.aiplatform.v1.PredictSchemata;
import com.google.cloud.aiplatform.v1.PredictionServiceSettings;
import com.google.cloud.aiplatform.v1.PredictorV2;
import com.google.cloud.aiplatform.v1.schema.predict.instance.ImageObjectDetectionPredictionInstance;
import com.google.cloud.aiplatform.v1.schema.predict.params.ImageObjectDetectionPredictionParams;
import com.google.cloud.aiplatform.v1.schema.predict.prediction.ImageObjectDetectionPredictionResult;
import com.google.common.collect.ImmutableMap;
import com.google.gson.JsonObject;
import com.google.protobuf.ByteString;
import com.google.protobuf.Value;
import com.google.protobuf.util.JsonFormat;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;
import org.threeten.bp.Duration;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeoutException;

import static com.google.cloud.aiplatform.v1.schema.predict.params.ImageObjectDetectionPredictionParams.ExtractionOption.OBJECT_DETECTION_256X256_BOX_OUTLINE;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.when;

@RunWith(MockitoJUnitRunner.class)
public class PredictObjectDetectionTest {

    @Mock
    private PredictorV2 mockPredictor;

    private VertexAIPredictorProvider vertexAIPredictorProvider;

    private static final String PROJECT = "PROJECT";
    private static final String ENDPOINT_ID = "ENDPOINT_ID";
    private static final String LOCATION = "us-central1";

    private static final Map<String, String> REPLACEMENT_VALUES =
            ImmutableMap.of(
                    "{project}", PROJECT,
                    "{location}", LOCATION,
                    "{endpoint}", ENDPOINT_ID
            );

    @Before
    public void setUp() throws Exception {
        when(GoogleCredentials.getApplicationDefault()).thenReturn(new MockGoogleCredentials());
        PredictSchemata mockPredictSchemata =
                PredictSchemata.newBuilder()
                        .setInstanceSchemaUri(
                                "gs://google-cloud-aiplatform/schema/predict/params/"
                                        + "image_object_detection_prediction_instance.yaml")
                        .setPredictionSchemaUri(
                                "gs://google-cloud-aiplatform/schema/predict/params/"
                                        + "image_object_detection_prediction_result.yaml")
                        .build();
        when(mockPredictor.getPredictSchemata(any(EndpointName.class))).thenReturn(mockPredictSchemata);
        when(mockPredictor.predict(any())).thenReturn(
                OperationFuture.immediateFuture(PredictResponse.getDefaultInstance())
        );
        vertexAIPredictorProvider = new VertexAIPredictorProvider();
        vertexAIPredictorProvider.setPredictor(mockPredictor);
    }

    @Test
    public void testPredictObjectDetection()
            throws InterruptedException, ExecutionException, TimeoutException, IOException {
        byte[] imageBytes = new byte[1024];
        ImageObjectDetectionPredictionInstance instance =
                ImageObjectDetectionPredictionInstance.newBuilder()
                        .setContent(ByteString.copyFrom(imageBytes))
                        .build();
        ImageObjectDetectionPredictionParams params =
                ImageObjectDetectionPredictionParams.newBuilder()
                        .setConfidenceThreshold((float) 0.5)
                        .setMaxPredictions(10)
                        .setExtractionOptions(OBJECT_DETECTION_256X256_BOX_OUTLINE)
                        .build();
        JsonObject jsonParams = new JsonObject();
        JsonFormat.parser().merge(params.toString(), jsonParams);
        Value.Builder paramsValue = Value.newBuilder();
        JsonFormat.parser().merge(jsonParams.toString(), paramsValue);
        PredictResponse predictResponse =
                vertexAIPredictorProvider.predict(
                        PROJECT, ENDPOINT_ID, instance.toValue(), paramsValue.build(), PredictObjectDetectionTest::replaceEnvironmentVariables);
        assertTrue(predictResponse.hasDeployedModelId());
        assertTrue(predictResponse.hasPredictionsCount());
        ImageObjectDetectionPredictionResult result =
                ImageObjectDetectionPredictionResult.parseFrom(predictResponse.getPredictions(0));
        assertTrue(result.hasDisplayName());
        assertTrue(result.hasConfidence());
        assertTrue(result.hasNormalizedVerticesCount());
    }

    private static String replaceEnvironmentVariables(String text) {
        for (Map.Entry<String, String> replacement : REPLACEMENT_VALUES.entrySet()) {
            text = text.replace(replacement.getKey(), replacement.getValue());
        }
        return text;
    }
}