|
7 | 7 | import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig; |
8 | 8 |
|
9 | 9 | import com.github.tomakehurst.wiremock.junit.WireMockRule; |
| 10 | +import com.github.tomakehurst.wiremock.stubbing.ServeEvent; |
| 11 | +import com.google.gson.Gson; |
| 12 | +import com.google.gson.GsonBuilder; |
| 13 | +import com.segment.analytics.gson.AutoValueAdapterFactory; |
| 14 | +import com.segment.analytics.gson.ISO8601DateAdapter; |
10 | 15 | import com.segment.analytics.messages.TrackMessage; |
11 | | -import java.util.UUID; |
| 16 | +import java.util.ArrayList; |
| 17 | +import java.util.Arrays; |
| 18 | +import java.util.Date; |
| 19 | +import java.util.HashSet; |
| 20 | +import java.util.Iterator; |
| 21 | +import java.util.List; |
| 22 | +import java.util.concurrent.TimeUnit; |
| 23 | +import org.awaitility.Awaitility; |
12 | 24 | import org.junit.Before; |
13 | 25 | import org.junit.Rule; |
14 | 26 | import org.junit.Test; |
| 27 | +import wiremock.com.fasterxml.jackson.core.JsonProcessingException; |
| 28 | +import wiremock.com.fasterxml.jackson.databind.JsonNode; |
| 29 | +import wiremock.com.fasterxml.jackson.databind.ObjectMapper; |
15 | 30 |
|
16 | 31 | public class SegmentTest { |
17 | 32 |
|
18 | 33 | @Rule |
19 | | - public WireMockRule wireMockRule = new WireMockRule(wireMockConfig().dynamicPort(), false); |
| 34 | + public WireMockRule wireMockRule = |
| 35 | + new WireMockRule(wireMockConfig().dynamicPort().gzipDisabled(true), false); |
20 | 36 |
|
21 | 37 | Analytics analytics; |
22 | 38 |
|
| 39 | + GsonBuilder gsonBuilder = new GsonBuilder() |
| 40 | + .registerTypeAdapterFactory(new AutoValueAdapterFactory()) |
| 41 | + .registerTypeAdapter(Date.class, new ISO8601DateAdapter()); |
| 42 | + |
| 43 | + Gson gson = gsonBuilder.create(); |
| 44 | + |
23 | 45 | @Before |
24 | 46 | public void confWireMock() { |
25 | 47 | stubFor(post(urlEqualTo("/v1/import/")).willReturn(okJson("{\"success\": \"true\"}"))); |
26 | 48 |
|
27 | 49 | analytics = Analytics.builder("write-key") |
28 | 50 | .endpoint(wireMockRule.baseUrl()) |
| 51 | + .flushInterval(1, TimeUnit.SECONDS) |
| 52 | + .queueCapacity(500) |
29 | 53 | // callback |
30 | 54 | // http client |
31 | 55 | .build(); |
32 | 56 | } |
33 | 57 |
|
34 | 58 | @Test |
35 | | - public void test() { |
36 | | - analytics.enqueue(TrackMessage.builder("my-track") |
37 | | - .messageId(UUID.randomUUID().toString()) |
38 | | - .userId("userId")); |
| 59 | + public void test() throws Throwable { |
| 60 | + analytics.enqueue(TrackMessage.builder("my-track").messageId("m1").userId("userId")); |
| 61 | + analytics.enqueue(TrackMessage.builder("my-track").messageId("m2").userId("userId")); |
| 62 | + |
| 63 | + Awaitility.await().until(() -> sentMessagesEqualsTo("m1", "m2")); |
| 64 | + } |
| 65 | + |
| 66 | + @Test |
| 67 | + public void testMore() throws Throwable { |
| 68 | + System.err.println("wm at " + wireMockRule.baseUrl()); |
| 69 | + int num = 100_000; |
| 70 | + String[] expectedIds = new String[num]; |
| 71 | + for (int i = 0; i < num; i++) { |
| 72 | + String id = "m" + i; |
| 73 | + expectedIds[i] = id; |
| 74 | + analytics.enqueue(TrackMessage.builder("my-track").messageId(id).userId("userId")); |
| 75 | + } |
| 76 | + |
| 77 | + Awaitility.await().atMost(1, TimeUnit.MINUTES).until(() -> sentMessagesEqualsTo(expectedIds)); |
| 78 | + } |
| 79 | + |
| 80 | + private static final ObjectMapper OM = new ObjectMapper(); |
| 81 | + |
| 82 | + private boolean sentMessagesEqualsTo(String... msgIds) { |
| 83 | + return new HashSet<>(sentMessages()).equals(new HashSet<>(Arrays.asList(msgIds))); |
| 84 | + } |
| 85 | + |
| 86 | + private List<String> sentMessages() { |
| 87 | + List<String> messageIds = new ArrayList<>(); |
| 88 | + for (ServeEvent event : wireMockRule.getAllServeEvents()) { |
| 89 | + JsonNode batch; |
| 90 | + try { |
| 91 | + JsonNode json = OM.readTree(event.getRequest().getBodyAsString()); |
| 92 | + batch = json.get("batch"); |
| 93 | + if (batch == null) { |
| 94 | + continue; |
| 95 | + } |
| 96 | + } catch (JsonProcessingException e) { |
| 97 | + continue; |
| 98 | + } |
| 99 | + Iterator<JsonNode> msgs = batch.elements(); |
| 100 | + while (msgs.hasNext()) { |
| 101 | + messageIds.add(msgs.next().get("messageId").asText()); |
| 102 | + } |
| 103 | + } |
| 104 | + return messageIds; |
39 | 105 | } |
40 | 106 | } |
0 commit comments