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 feast.serving.service;
  
  import static feast.serving.util.DataGenerator.*;
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.hamcrest.Matchers.equalTo;
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.Mockito.when;
  import static org.mockito.MockitoAnnotations.initMocks;
  
  import com.google.common.collect.ImmutableMap;
  import com.google.protobuf.Duration;
  import com.google.protobuf.Timestamp;
  import feast.proto.core.FeatureProto;
  import feast.proto.core.FeatureViewProto;
  import feast.proto.serving.ServingAPIProto;
  import feast.proto.serving.ServingAPIProto.FieldStatus;
  import feast.proto.serving.ServingAPIProto.GetOnlineFeaturesResponse;
  import feast.proto.types.ValueProto;
  import feast.serving.connectors.Feature;
  import feast.serving.connectors.ProtoFeature;
  import feast.serving.connectors.redis.retriever.RedisOnlineRetriever;
  import feast.serving.registry.Registry;
  import feast.serving.registry.RegistryRepository;
  import io.opentracing.Tracer;
  import io.opentracing.Tracer.SpanBuilder;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.List;
  import java.util.Optional;
  import java.util.stream.Collectors;
  import org.junit.Before;
  import org.junit.Test;
  import org.mockito.ArgumentMatchers;
  import org.mockito.Mock;
  import org.mockito.Mockito;
  
  public class OnlineServingServiceTest {
  
    @Mock Registry registry;
    @Mock Tracer tracer;
    @Mock RedisOnlineRetriever retrieverV2;
    private String transformationServiceEndpoint;
  
    private OnlineServingServiceV2 onlineServingServiceV2;
  
    List<Feature> mockedFeatureRows;
    List<FeatureProto.FeatureSpecV2> featureSpecs;
  
    Timestamp now = Timestamp.newBuilder().setSeconds(System.currentTimeMillis() / 1000).build();
  
    @Before
    public void setUp() {
      initMocks(this);
  
      RegistryRepository registryRepo = new RegistryRepository(registry);
  
      OnlineTransformationService onlineTransformationService =
          new OnlineTransformationService(transformationServiceEndpoint, registryRepo);
      onlineServingServiceV2 =
          new OnlineServingServiceV2(
              retrieverV2,
              registryRepo,
              onlineTransformationService,
              "feast_project",
              Optional.of(tracer));
  
      mockedFeatureRows = new ArrayList<>();
      mockedFeatureRows.add(
          new ProtoFeature(
              ServingAPIProto.FeatureReferenceV2.newBuilder()
                  .setFeatureViewName("featureview_1")
                  .setFeatureName("feature_1")
                  .build(),
              now,
              createStrValue("1")));
      mockedFeatureRows.add(
          new ProtoFeature(
              ServingAPIProto.FeatureReferenceV2.newBuilder()
                  .setFeatureViewName("featureview_1")
                  .setFeatureName("feature_2")
                  .build(),
              now,
              createStrValue("2")));
      mockedFeatureRows.add(
          new ProtoFeature(
              ServingAPIProto.FeatureReferenceV2.newBuilder()
                  .setFeatureViewName("featureview_1")
                  .setFeatureName("feature_1")
                  .build(),
              now,
              createStrValue("3")));
      mockedFeatureRows.add(
          new ProtoFeature(
              ServingAPIProto.FeatureReferenceV2.newBuilder()
                  .setFeatureViewName("featureview_1")
                  .setFeatureName("feature_2")
                  .build(),
              now,
              createStrValue("4")));
      mockedFeatureRows.add(
          new ProtoFeature(
              ServingAPIProto.FeatureReferenceV2.newBuilder()
                  .setFeatureViewName("featureview_1")
                  .setFeatureName("feature_3")
                  .build(),
              now,
              createStrValue("5")));
      mockedFeatureRows.add(
          new ProtoFeature(
              ServingAPIProto.FeatureReferenceV2.newBuilder()
                  .setFeatureViewName("featureview_1")
                  .setFeatureName("feature_1")
                  .build(),
              Timestamp.newBuilder().setSeconds(1).build(),
              createStrValue("6")));
  
      featureSpecs = new ArrayList<>();
      featureSpecs.add(
          FeatureProto.FeatureSpecV2.newBuilder()
              .setName("feature_1")
              .setValueType(ValueProto.ValueType.Enum.STRING)
              .build());
      featureSpecs.add(
          FeatureProto.FeatureSpecV2.newBuilder()
              .setName("feature_2")
              .setValueType(ValueProto.ValueType.Enum.STRING)
              .build());
    }
  
    @Test
    public void shouldReturnResponseWithValuesAndMetadataIfKeysPresent() {
      String projectName = "default";
      ServingAPIProto.FeatureReferenceV2 featureReference1 =
          ServingAPIProto.FeatureReferenceV2.newBuilder()
              .setFeatureViewName("featureview_1")
              .setFeatureName("feature_1")
              .build();
      ServingAPIProto.FeatureReferenceV2 featureReference2 =
          ServingAPIProto.FeatureReferenceV2.newBuilder()
              .setFeatureViewName("featureview_1")
              .setFeatureName("feature_2")
              .build();
      List<ServingAPIProto.FeatureReferenceV2> featureReferences =
          List.of(featureReference1, featureReference2);
      ServingAPIProto.GetOnlineFeaturesRequest request = getOnlineFeaturesRequest(featureReferences);
  
      List<List<Feature>> featureRows =
          List.of(
              List.of(mockedFeatureRows.get(0), mockedFeatureRows.get(1)),
              List.of(mockedFeatureRows.get(2), mockedFeatureRows.get(3)));
  
      when(retrieverV2.getOnlineFeatures(any(), any(), any())).thenReturn(featureRows);
      when(registry.getFeatureViewSpec(any())).thenReturn(getFeatureViewSpec());
      when(registry.getFeatureSpec(mockedFeatureRows.get(0).getFeatureReference()))
          .thenReturn(featureSpecs.get(0));
      when(registry.getFeatureSpec(mockedFeatureRows.get(1).getFeatureReference()))
          .thenReturn(featureSpecs.get(1));
      when(registry.getFeatureSpec(mockedFeatureRows.get(2).getFeatureReference()))
          .thenReturn(featureSpecs.get(0));
      when(registry.getFeatureSpec(mockedFeatureRows.get(3).getFeatureReference()))
          .thenReturn(featureSpecs.get(1));
      when(registry.getEntityJoinKey("entity1")).thenReturn("entity1");
      when(registry.getEntityJoinKey("entity2")).thenReturn("entity2");
  
      when(tracer.buildSpan(ArgumentMatchers.any())).thenReturn(Mockito.mock(SpanBuilder.class));
  
      GetOnlineFeaturesResponse expected =
          GetOnlineFeaturesResponse.newBuilder()
              .addResults(
                  GetOnlineFeaturesResponse.FeatureVector.newBuilder()
                      .addValues(createStrValue("1"))
                      .addValues(createStrValue("3"))
                      .addStatuses(FieldStatus.PRESENT)
                      .addStatuses(FieldStatus.PRESENT)
                      .addEventTimestamps(now)
                      .addEventTimestamps(now))
              .addResults(
                  GetOnlineFeaturesResponse.FeatureVector.newBuilder()
                      .addValues(createStrValue("2"))
                      .addValues(createStrValue("4"))
                      .addStatuses(FieldStatus.PRESENT)
                      .addStatuses(FieldStatus.PRESENT)
                      .addEventTimestamps(now)
                      .addEventTimestamps(now))
              .setMetadata(
                  ServingAPIProto.GetOnlineFeaturesResponseMetadata.newBuilder()
                      .setFeatureNames(
                          ServingAPIProto.FeatureList.newBuilder()
                              .addVal("featureview_1:feature_1")
                              .addVal("featureview_1:feature_2")))
              .build();
      ServingAPIProto.GetOnlineFeaturesResponse actual =
          onlineServingServiceV2.getOnlineFeatures(request);
      assertThat(actual, equalTo(expected));
    }
  
    @Test
    public void shouldReturnResponseWithUnsetValuesAndMetadataIfKeysNotPresent() {
      String projectName = "default";
      ServingAPIProto.FeatureReferenceV2 featureReference1 =
          ServingAPIProto.FeatureReferenceV2.newBuilder()
              .setFeatureViewName("featureview_1")
              .setFeatureName("feature_1")
              .build();
      ServingAPIProto.FeatureReferenceV2 featureReference2 =
          ServingAPIProto.FeatureReferenceV2.newBuilder()
              .setFeatureViewName("featureview_1")
              .setFeatureName("feature_2")
              .build();
      List<ServingAPIProto.FeatureReferenceV2> featureReferences =
          List.of(featureReference1, featureReference2);
      ServingAPIProto.GetOnlineFeaturesRequest request = getOnlineFeaturesRequest(featureReferences);
  
      List<Feature> entityKeyList1 = new ArrayList<>();
      List<Feature> entityKeyList2 = new ArrayList<>();
      entityKeyList1.add(mockedFeatureRows.get(0));
      entityKeyList1.add(mockedFeatureRows.get(1));
      entityKeyList2.add(mockedFeatureRows.get(4));
  
      List<List<Feature>> featureRows =
          List.of(
              List.of(mockedFeatureRows.get(0), mockedFeatureRows.get(1)),
              Arrays.asList(null, mockedFeatureRows.get(4)));
  
      when(retrieverV2.getOnlineFeatures(any(), any(), any())).thenReturn(featureRows);
      when(registry.getFeatureViewSpec(any())).thenReturn(getFeatureViewSpec());
      when(registry.getFeatureSpec(mockedFeatureRows.get(0).getFeatureReference()))
          .thenReturn(featureSpecs.get(0));
      when(registry.getFeatureSpec(mockedFeatureRows.get(1).getFeatureReference()))
          .thenReturn(featureSpecs.get(1));
      when(registry.getEntityJoinKey("entity1")).thenReturn("entity1");
      when(registry.getEntityJoinKey("entity2")).thenReturn("entity2");
  
      when(tracer.buildSpan(ArgumentMatchers.any())).thenReturn(Mockito.mock(SpanBuilder.class));
  
      GetOnlineFeaturesResponse expected =
          GetOnlineFeaturesResponse.newBuilder()
              .addResults(
                  GetOnlineFeaturesResponse.FeatureVector.newBuilder()
                      .addValues(createStrValue("1"))
                      .addValues(createEmptyValue())
                      .addStatuses(FieldStatus.PRESENT)
                      .addStatuses(FieldStatus.NOT_FOUND)
                      .addEventTimestamps(now)
                      .addEventTimestamps(Timestamp.newBuilder().build()))
              .addResults(
                  GetOnlineFeaturesResponse.FeatureVector.newBuilder()
                      .addValues(createStrValue("2"))
                      .addValues(createStrValue("5"))
                      .addStatuses(FieldStatus.PRESENT)
                      .addStatuses(FieldStatus.PRESENT)
                      .addEventTimestamps(now)
                      .addEventTimestamps(now))
              .setMetadata(
                  ServingAPIProto.GetOnlineFeaturesResponseMetadata.newBuilder()
                      .setFeatureNames(
                          ServingAPIProto.FeatureList.newBuilder()
                              .addVal("featureview_1:feature_1")
                              .addVal("featureview_1:feature_2")))
              .build();
      GetOnlineFeaturesResponse actual = onlineServingServiceV2.getOnlineFeatures(request);
      assertThat(actual, equalTo(expected));
    }
  
    @Test
    public void shouldReturnResponseWithValuesAndMetadataIfMaxAgeIsExceeded() {
      String projectName = "default";
      ServingAPIProto.FeatureReferenceV2 featureReference1 =
          ServingAPIProto.FeatureReferenceV2.newBuilder()
              .setFeatureViewName("featureview_1")
              .setFeatureName("feature_1")
              .build();
      ServingAPIProto.FeatureReferenceV2 featureReference2 =
          ServingAPIProto.FeatureReferenceV2.newBuilder()
              .setFeatureViewName("featureview_1")
              .setFeatureName("feature_2")
              .build();
      List<ServingAPIProto.FeatureReferenceV2> featureReferences =
          List.of(featureReference1, featureReference2);
      ServingAPIProto.GetOnlineFeaturesRequest request = getOnlineFeaturesRequest(featureReferences);
  
      List<List<Feature>> featureRows =
          List.of(
              List.of(mockedFeatureRows.get(5), mockedFeatureRows.get(1)),
              List.of(mockedFeatureRows.get(5), mockedFeatureRows.get(1)));
  
      when(retrieverV2.getOnlineFeatures(any(), any(), any())).thenReturn(featureRows);
      when(registry.getFeatureViewSpec(any()))
          .thenReturn(
              FeatureViewProto.FeatureViewSpec.newBuilder()
                  .setName("featureview_1")
                  .addEntities("entity1")
                  .addEntities("entity2")
                  .addFeatures(
                      FeatureProto.FeatureSpecV2.newBuilder()
                          .setName("feature_1")
                          .setValueType(ValueProto.ValueType.Enum.STRING)
                          .build())
                  .addFeatures(
                      FeatureProto.FeatureSpecV2.newBuilder()
                          .setName("feature_2")
                          .setValueType(ValueProto.ValueType.Enum.STRING)
                          .build())
                  .setTtl(Duration.newBuilder().setSeconds(3600))
                  .build());
      when(registry.getFeatureSpec(mockedFeatureRows.get(1).getFeatureReference()))
          .thenReturn(featureSpecs.get(1));
      when(registry.getFeatureSpec(mockedFeatureRows.get(5).getFeatureReference()))
          .thenReturn(featureSpecs.get(0));
      when(registry.getEntityJoinKey("entity1")).thenReturn("entity1");
      when(registry.getEntityJoinKey("entity2")).thenReturn("entity2");
  
      when(tracer.buildSpan(ArgumentMatchers.any())).thenReturn(Mockito.mock(SpanBuilder.class));
  
      GetOnlineFeaturesResponse expected =
          GetOnlineFeaturesResponse.newBuilder()
              .addResults(
                  GetOnlineFeaturesResponse.FeatureVector.newBuilder()
                      .addValues(createStrValue("6"))
                      .addValues(createStrValue("6"))
                      .addStatuses(FieldStatus.OUTSIDE_MAX_AGE)
                      .addStatuses(FieldStatus.OUTSIDE_MAX_AGE)
                      .addEventTimestamps(Timestamp.newBuilder().setSeconds(1).build())
                      .addEventTimestamps(Timestamp.newBuilder().setSeconds(1).build()))
              .addResults(
                  GetOnlineFeaturesResponse.FeatureVector.newBuilder()
                      .addValues(createStrValue("2"))
                      .addValues(createStrValue("2"))
                      .addStatuses(FieldStatus.PRESENT)
                      .addStatuses(FieldStatus.PRESENT)
                      .addEventTimestamps(now)
                      .addEventTimestamps(now))
              .setMetadata(
                  ServingAPIProto.GetOnlineFeaturesResponseMetadata.newBuilder()
                      .setFeatureNames(
                          ServingAPIProto.FeatureList.newBuilder()
                              .addVal("featureview_1:feature_1")
                              .addVal("featureview_1:feature_2")))
              .build();
      GetOnlineFeaturesResponse actual = onlineServingServiceV2.getOnlineFeatures(request);
      assertThat(actual, equalTo(expected));
    }
  
    private FeatureViewProto.FeatureViewSpec getFeatureViewSpec() {
      return FeatureViewProto.FeatureViewSpec.newBuilder()
          .setName("featureview_1")
          .addEntities("entity1")
          .addEntities("entity2")
          .addFeatures(
              FeatureProto.FeatureSpecV2.newBuilder()
                  .setName("feature_1")
                  .setValueType(ValueProto.ValueType.Enum.STRING)
                  .build())
          .addFeatures(
              FeatureProto.FeatureSpecV2.newBuilder()
                  .setName("feature_2")
                  .setValueType(ValueProto.ValueType.Enum.STRING)
                  .build())
          .setTtl(Duration.newBuilder().setSeconds(120))
          .build();
    }
  
    private ServingAPIProto.GetOnlineFeaturesRequest getOnlineFeaturesRequest(
        List<ServingAPIProto.FeatureReferenceV2> featureReferences) {
      return ServingAPIProto.GetOnlineFeaturesRequest.newBuilder()
          .setFeatures(
              ServingAPIProto.FeatureList.newBuilder()
                  .addAllVal(
                      featureReferences.stream()
                          .map(FeatureUtil::getFeatureReference)
                          .collect(Collectors.toList()))
                  .build())
          .putAllEntities(
              ImmutableMap.of(
                  "entity1",
                      ValueProto.RepeatedValue.newBuilder()
                          .addAllVal(List.of(createInt64Value(1), createInt64Value(2)))
                          .build(),
                  "entity2",
                      ValueProto.RepeatedValue.newBuilder()
                          .addAllVal(List.of(createStrValue("a"), createStrValue("b")))
                          .build()))
          .build();
    }
  }
  