GrpcMetricsContentTypeTest.java

/*
 * Copyright 2017 LINE Corporation
 *
 * LINE Corporation licenses this file to you 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 com.linecorp.armeria.it.grpc;

import static io.micrometer.core.instrument.Statistic.COUNT;
import static org.assertj.core.api.Assertions.assertThat;
import static org.awaitility.Awaitility.given;

import java.util.Optional;
import java.util.concurrent.TimeUnit;

import org.junit.AfterClass;
import org.junit.ClassRule;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.DisableOnDebug;
import org.junit.rules.TestRule;
import org.junit.rules.Timeout;

import com.google.protobuf.ByteString;

import com.linecorp.armeria.client.ClientBuilder;
import com.linecorp.armeria.client.ClientFactory;
import com.linecorp.armeria.client.ClientFactoryBuilder;
import com.linecorp.armeria.client.HttpClient;
import com.linecorp.armeria.common.AggregatedHttpMessage;
import com.linecorp.armeria.common.HttpHeaderNames;
import com.linecorp.armeria.common.HttpHeaders;
import com.linecorp.armeria.common.HttpMethod;
import com.linecorp.armeria.common.SessionProtocol;
import com.linecorp.armeria.common.grpc.GrpcSerializationFormats;
import com.linecorp.armeria.common.metric.MeterIdPrefix;
import com.linecorp.armeria.common.metric.MeterIdPrefixFunction;
import com.linecorp.armeria.common.metric.MoreMeters;
import com.linecorp.armeria.common.metric.PrometheusMeterRegistries;
import com.linecorp.armeria.grpc.testing.Messages;
import com.linecorp.armeria.grpc.testing.TestServiceGrpc;
import com.linecorp.armeria.server.ServerBuilder;
import com.linecorp.armeria.server.grpc.GrpcServiceBuilder;
import com.linecorp.armeria.server.metric.MetricCollectingService;
import com.linecorp.armeria.testing.server.ServerRule;

import io.grpc.stub.StreamObserver;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Statistic;

public class GrpcMetricsContentTypeTest {
    private static final MeterRegistry registry = PrometheusMeterRegistries.newRegistry();

    private static final Messages.SimpleResponse OK_RESPONSE = Messages.SimpleResponse.newBuilder().setUsername(
            "world").build();

    private static class TestServiceImpl extends TestServiceGrpc.TestServiceImplBase {
        @Override
        public void unaryCall(Messages.SimpleRequest request,
                              StreamObserver<Messages.SimpleResponse> responseObserver) {
            if ("world".equals(request.getPayload().getBody().toStringUtf8())) {
                responseObserver.onNext(OK_RESPONSE);
                responseObserver.onCompleted();
                return;
            }
            responseObserver.onError(new IllegalArgumentException("bad argument"));
        }
    }

    @ClassRule
    public static final ServerRule server = new ServerRule() {
        @Override
        protected void configure(ServerBuilder sb) throws Exception {
            sb.meterRegistry(registry);
            sb.port(0, SessionProtocol.HTTP);
            sb.serviceUnder("/", new GrpcServiceBuilder()
                    .addService(new GrpcMetricsContentTypeTest.TestServiceImpl())
                    .enableUnframedRequests(true)
                    .build()
                    .decorate(MetricCollectingService.newDecorator(
                            MeterIdPrefixFunction.ofDefault("server"))));
        }
    };

    private static final ClientFactory clientFactory =
            new ClientFactoryBuilder().build();

    @AfterClass
    public static void closeClientFactory() {
        clientFactory.close();
    }

    @Rule
    public final TestRule globalTimeout = new DisableOnDebug(new Timeout(30, TimeUnit.SECONDS));

    @Test
    public void normalGrpc() throws Exception {
        // makeRequest("world"); // test passed
        makeUnframedRequest("world"); // test failed

        given().ignoreExceptions().untilAsserted(() -> assertThat(
                findServerMeter("UnaryCall", "requests", COUNT, "result", "success")).contains(1.0));
        given().ignoreExceptions().untilAsserted(() -> assertThat(
                findServerMeter("UnaryCall", "requests", COUNT, "result", "failure")).contains(0.0));
    }

    private static Optional<Double> findServerMeter(
            String method, String suffix, Statistic type, String... keyValues) {
        final MeterIdPrefix prefix = new MeterIdPrefix(
                "server." + suffix + '#' + type.getTagValueRepresentation(),
                "method", "armeria.grpc.testing.TestService/" + method,
                "hostnamePattern", "*",
                "pathMapping", "catch-all");
        final String meterIdStr = prefix.withTags(keyValues).toString();
        return Optional.ofNullable(MoreMeters.measureAll(registry).get(meterIdStr));
    }

    private static void makeRequest(String name) throws Exception {
        TestServiceGrpc.TestServiceBlockingStub client = new ClientBuilder(
                server.uri(GrpcSerializationFormats.PROTO, "/"))
                .factory(clientFactory)
                .build(TestServiceGrpc.TestServiceBlockingStub.class);

        Messages.SimpleRequest request =
                Messages.SimpleRequest.newBuilder()
                                      .setPayload(Messages.Payload.newBuilder()
                                                                  .setBody(ByteString.copyFromUtf8(name)))
                                      .build();
        try {
            client.unaryCall(request);
        } catch (Throwable t) {
            // Ignore, we will count these up
        }
    }

    private static void makeUnframedRequest(String name) throws Exception {
        Messages.SimpleRequest msg =
                Messages.SimpleRequest.newBuilder()
                                      .setPayload(Messages.Payload.newBuilder()
                                                                  .setBody(ByteString.copyFromUtf8(name)))
                                      .build();

        HttpClient client = HttpClient.of(server.httpUri("/"));
        AggregatedHttpMessage response = client.execute(
                HttpHeaders.of(HttpMethod.POST,
                               TestServiceGrpc.getUnaryCallMethod().getFullMethodName())
                           .set(HttpHeaderNames.CONTENT_TYPE, "application/protobuf"),
                msg.toByteArray()).aggregate().get();

        Messages.SimpleResponse message = Messages.SimpleResponse.parseFrom(response.content().array());
        assertThat(message).isEqualTo(OK_RESPONSE);
    }
}