package com.stripe.functional;

import static junit.framework.TestCase.assertEquals;
import static junit.framework.TestCase.assertNull;
import static junit.framework.TestCase.assertTrue;

import com.stripe.BaseStripeTest;
import com.stripe.Stripe;
import com.stripe.exception.StripeException;
import com.stripe.model.Balance;
import com.stripe.net.ApiResource;
import com.stripe.net.ClientTelemetryPayload;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.TimeUnit;

import okhttp3.mockwebserver.MockResponse;
import okhttp3.mockwebserver.MockWebServer;
import okhttp3.mockwebserver.RecordedRequest;
import org.junit.Test;


public class TelemetryTest extends BaseStripeTest {
  @Test
  public void testTelemetryEnabled() throws StripeException, IOException, InterruptedException {
    MockWebServer server = new MockWebServer();
    server.enqueue(new MockResponse().setBody("{}").addHeader("Request-Id", "req_1")
        .setBodyDelay(30, TimeUnit.MILLISECONDS));
    server.enqueue(new MockResponse().setBody("{}").addHeader("Request-Id", "req_2")
        .setBodyDelay(120, TimeUnit.MILLISECONDS));
    server.enqueue(new MockResponse().setBody("{}").addHeader("Request-Id", "req_3"));
    server.start();

    Stripe.overrideApiBase(server.url("").toString());
    Stripe.enableTelemetry = true;

    Balance b1 = Balance.retrieve();
    RecordedRequest request1 = server.takeRequest();
    assertNull(request1.getHeader("X-Stripe-Client-Telemetry"));

    Balance b2 = Balance.retrieve();
    RecordedRequest request2 = server.takeRequest();
    String telemetry1 = request2.getHeader("X-Stripe-Client-Telemetry");
    ClientTelemetryPayload payload1 = ApiResource.GSON.fromJson(
        telemetry1, ClientTelemetryPayload.class);
    assertEquals(payload1.lastRequestMetrics.requestId, "req_1");
    assertTrue(payload1.lastRequestMetrics.requestDurationMs > 30);
    assertTrue(payload1.lastRequestMetrics.requestDurationMs < 130);

    Balance b3 = Balance.retrieve();
    RecordedRequest request3 = server.takeRequest();
    String telemetry2 = request3.getHeader("X-Stripe-Client-Telemetry");
    ClientTelemetryPayload payload2 = ApiResource.GSON.fromJson(
        telemetry2, ClientTelemetryPayload.class);
    assertEquals(payload2.lastRequestMetrics.requestId, "req_2");
    assertTrue(payload2.lastRequestMetrics.requestDurationMs > 120);
    assertTrue(payload2.lastRequestMetrics.requestDurationMs < 220);

    server.shutdown();
  }

  @Test
  public void testTelemetryDisabled() throws StripeException, IOException, InterruptedException {
    MockWebServer server = new MockWebServer();
    server.enqueue(new MockResponse().setBody("{}").addHeader("Request-Id", "req_1"));
    server.enqueue(new MockResponse().setBody("{}").addHeader("Request-Id", "req_2"));
    server.enqueue(new MockResponse().setBody("{}").addHeader("Request-Id", "req_3"));
    server.start();

    Stripe.overrideApiBase(server.url("").toString());
    Stripe.enableTelemetry = false;

    Balance b1 = Balance.retrieve();
    RecordedRequest request1 = server.takeRequest();
    assertNull(request1.getHeader("X-Stripe-Client-Telemetry"));

    Balance b2 = Balance.retrieve();
    RecordedRequest request2 = server.takeRequest();
    assertNull(request2.getHeader("X-Stripe-Client-Telemetry"));

    server.shutdown();
  }

  @Test
  public void testTelemetryWorksWithConcurrentRequests() throws IOException, InterruptedException {
    MockWebServer server = new MockWebServer();

    for (int i = 0; i < 20; i++) {
      server.enqueue(new MockResponse().setBody("{}")
          .addHeader("Request-Id", "req_" + i)
          .setBodyDelay(50, TimeUnit.MILLISECONDS));
    }
    server.start();

    Stripe.overrideApiBase(server.url("").toString());
    Stripe.enableTelemetry = true;

    Runnable work = new Runnable() {
      @Override
      public void run() {
        try {
          Balance.retrieve();
        } catch (StripeException e) {
          assertNull(e);
        }
      }
    };

    // the first 10 requests will not contain telemetry
    ArrayList<Thread> threads = new ArrayList<>();
    for (int i = 0; i < 10; i++) {
      threads.add(new Thread(work));
    }
    for (int i = 0; i < 10; i++) {
      threads.get(i).start();
    }
    for (int i = 0; i < 10; i++) {
      threads.get(i).join();
    }
    threads.clear();

    // the following 10 requests will contain telemetry
    for (int i = 0; i < 10; i++) {
      threads.add(new Thread(work));
    }
    for (int i = 0; i < 10; i++) {
      threads.get(i).start();
    }
    for (int i = 0; i < 10; i++) {
      threads.get(i).join();
    }
    threads.clear();

    Set<String> seenRequestIds = new HashSet<>();

    for (int i = 0; i < 10; i++) {
      RecordedRequest request = server.takeRequest();
      assertNull(request.getHeader("X-Stripe-Client-Telemetry"));
    }

    for (int i = 0; i < 10; i++) {
      RecordedRequest request = server.takeRequest();
      String telemetry2 = request.getHeader("X-Stripe-Client-Telemetry");
      ClientTelemetryPayload payload = ApiResource.GSON.fromJson(
          telemetry2, ClientTelemetryPayload.class);
      seenRequestIds.add(payload.lastRequestMetrics.requestId);
    }

    // check that each telemetry payload corresponds to a unique request id
    assertEquals(10, seenRequestIds.size());

    server.shutdown();
  }
}

@@ -1951,77 +1951,8 @@
 30);
-%0A    assertTrue(payload1.lastRequestMetrics.requestDurationMs %3C 130);
 %0A%0A  
@@ -2366,77 +2366,8 @@
 20);
-%0A    assertTrue(payload2.lastRequestMetrics.requestDurationMs %3C 220);
 %0A%0A  

Remove upper bounds in timing assertions because slowness is inevitable