package com.salesforce.reactorgrpc;

import com.salesforce.grpc.testing.contrib.NettyGrpcServerRule;
import org.junit.Rule;
import org.junit.Test;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.test.StepVerifier;
import reactor.util.context.Context;

import java.util.stream.Collectors;

import static org.assertj.core.api.Assertions.assertThat;

public class ReactorContextPropagationTest {

    @Rule
    public NettyGrpcServerRule serverRule = new NettyGrpcServerRule();

    private static class SimpleGreeter extends ReactorGreeterGrpc.GreeterImplBase {
        @Override
        public Mono<HelloResponse> sayHello(Mono<HelloRequest> request) {
            return request.map(HelloRequest::getName)
                    .map(name -> HelloResponse.newBuilder().setMessage("Hello " + name).build());
        }

        @Override
        public Mono<HelloResponse> sayHelloReqStream(Flux<HelloRequest> request) {
            return request.map(HelloRequest::getName)
                    .collect(Collectors.joining("and"))
                    .map(names -> HelloResponse.newBuilder().setMessage("Hello " + names).build());
        }
    }

    @Test
    public void contextDoesNotPropagate() {
        serverRule.getServiceRegistry().addService(new SimpleGreeter());

        ReactorGreeterGrpc.ReactorGreeterStub stub = ReactorGreeterGrpc.newReactorStub(serverRule.getChannel());
        Mono<HelloRequest> req = Mono.just(HelloRequest.newBuilder().setName("reactor").build());

        Mono<HelloResponse> resp = req
                // This assertion will fail 
                .doOnEach(signal -> assertThat(signal.getContext().getOrEmpty("key")).isNotEmpty())
                .compose(stub::sayHello)
                // This assertion won't fail
                .doOnEach(signal -> assertThat(signal.getContext().getOrEmpty("key")).isNotEmpty())
                .subscriberContext(Context.of("key", "dummy"));

        StepVerifier.create(resp)
                .expectNextMatches(response -> response.getMessage().equals("Hello reactor"))
                .verifyComplete();
    }
}