code: /*
   * SPDX-License-Identifier: Apache-2.0
   * Copyright 2018-2019 The Feast Authors
   *
   * 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
   *
   *     https://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 dev.feast;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.mockito.AdditionalAnswers.delegatesTo;
  import static org.mockito.Mockito.mock;
  
  import com.google.protobuf.Timestamp;
  import feast.proto.serving.ServingAPIProto;
  import feast.proto.serving.ServingAPIProto.FieldStatus;
  import feast.proto.serving.ServingAPIProto.GetOnlineFeaturesRequest;
  import feast.proto.serving.ServingAPIProto.GetOnlineFeaturesResponse;
  import feast.proto.serving.ServingServiceGrpc.ServingServiceImplBase;
  import feast.proto.types.ValueProto;
  import feast.proto.types.ValueProto.Value;
  import io.grpc.*;
  import io.grpc.inprocess.InProcessChannelBuilder;
  import io.grpc.inprocess.InProcessServerBuilder;
  import io.grpc.stub.StreamObserver;
  import io.grpc.testing.GrpcCleanupRule;
  import java.time.Instant;
  import java.util.Arrays;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Optional;
  import java.util.concurrent.atomic.AtomicBoolean;
  import org.junit.Before;
  import org.junit.Rule;
  import org.junit.Test;
  
  public class FeastClientTest {
    private final String AUTH_TOKEN = "test token";
  
    @Rule public GrpcCleanupRule grpcRule;
    private AtomicBoolean isAuthenticated;
  
    private ServingServiceImplBase servingMock =
        mock(
            ServingServiceImplBase.class,
            delegatesTo(
                new ServingServiceImplBase() {
                  @Override
                  public void getOnlineFeatures(
                      GetOnlineFeaturesRequest request,
                      StreamObserver<GetOnlineFeaturesResponse> responseObserver) {
                    if (!request.equals(FeastClientTest.getFakeRequest())) {
                      responseObserver.onError(Status.FAILED_PRECONDITION.asRuntimeException());
                    }
  
                    responseObserver.onNext(FeastClientTest.getFakeResponse());
                    responseObserver.onCompleted();
                  }
                }));
  
    private FeastClient client;
  
    @Before
    public void setup() throws Exception {
      this.grpcRule = new GrpcCleanupRule();
      this.isAuthenticated = new AtomicBoolean(false);
      // setup fake serving service
      String serverName = InProcessServerBuilder.generateName();
      this.grpcRule.register(
          InProcessServerBuilder.forName(serverName)
              .directExecutor()
              .addService(this.servingMock)
              .build()
              .start());
  
      // setup test feast client target
      ManagedChannel channel =
          this.grpcRule.register(
              InProcessChannelBuilder.forName(serverName).directExecutor().build());
      this.client = new FeastClient(channel, Optional.empty());
    }
  
    @Test
    public void shouldGetOnlineFeatures() {
      shouldGetOnlineFeaturesWithClient(this.client);
    }
  
    private void shouldGetOnlineFeaturesWithClient(FeastClient client) {
      List<Row> rows =
          client.getOnlineFeatures(
              Arrays.asList("driver:name", "driver:rating", "driver:null_value"),
              Arrays.asList(
                  Row.create().set("driver_id", 1).setEntityTimestamp(Instant.ofEpochSecond(100))),
              "driver_project");
  
      assertEquals(
          rows.get(0).getFields(),
          new HashMap<String, Value>() {
            {
              put("driver_id", intValue(1));
              put("driver:name", strValue("david"));
              put("driver:rating", intValue(3));
              put("driver:null_value", Value.newBuilder().build());
            }
          });
      assertEquals(
          rows.get(0).getStatuses(),
          new HashMap<String, FieldStatus>() {
            {
              put("driver_id", FieldStatus.PRESENT);
              put("driver:name", FieldStatus.PRESENT);
              put("driver:rating", FieldStatus.PRESENT);
              put("driver:null_value", FieldStatus.NULL_VALUE);
            }
          });
    }
  
    private static GetOnlineFeaturesRequest getFakeRequest() {
      // setup mock serving service stub
      return GetOnlineFeaturesRequest.newBuilder()
          .setFeatures(
              ServingAPIProto.FeatureList.newBuilder()
                  .addVal("driver:name")
                  .addVal("driver:rating")
                  .addVal("driver:null_value")
                  .build())
          .putEntities("driver_id", ValueProto.RepeatedValue.newBuilder().addVal(intValue(1)).build())
          .build();
    }
  
    private static GetOnlineFeaturesResponse getFakeResponse() {
      return GetOnlineFeaturesResponse.newBuilder()
          .addResults(
              GetOnlineFeaturesResponse.FeatureVector.newBuilder()
                  .addValues(strValue("david"))
                  .addStatuses(FieldStatus.PRESENT)
                  .addEventTimestamps(Timestamp.newBuilder())
                  .build())
          .addResults(
              GetOnlineFeaturesResponse.FeatureVector.newBuilder()
                  .addValues(intValue(3))
                  .addStatuses(FieldStatus.PRESENT)
                  .addEventTimestamps(Timestamp.newBuilder())
                  .build())
          .addResults(
              GetOnlineFeaturesResponse.FeatureVector.newBuilder()
                  .addValues(Value.newBuilder().build())
                  .addStatuses(FieldStatus.NULL_VALUE)
                  .addEventTimestamps(Timestamp.newBuilder())
                  .build())
          .setMetadata(
              ServingAPIProto.GetOnlineFeaturesResponseMetadata.newBuilder()
                  .setFeatureNames(
                      ServingAPIProto.FeatureList.newBuilder()
                          .addVal("driver:name")
                          .addVal("driver:rating")
                          .addVal("driver:null_value"))
                  .build())
          .build();
    }
  
    private static Value strValue(String val) {
      return Value.newBuilder().setStringVal(val).build();
    }
  
    private static Value intValue(int val) {
      return Value.newBuilder().setInt32Val(val).build();
    }
  }
  