/*

* 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();

}

}