1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28 package org.apache.hc.client5.http.impl.classic;
29
30 import org.apache.hc.client5.http.classic.ExecRuntime;
31 import org.apache.hc.core5.function.Supplier;
32 import org.apache.hc.core5.http.ClassicHttpResponse;
33 import org.apache.hc.core5.http.Header;
34 import org.apache.hc.core5.http.HttpEntity;
35 import org.apache.hc.core5.http.impl.io.ChunkedInputStream;
36 import org.apache.hc.core5.http.impl.io.SessionInputBufferImpl;
37 import org.apache.hc.core5.http.io.SessionInputBuffer;
38 import org.junit.Assert;
39 import org.junit.Before;
40 import org.junit.Test;
41 import org.mockito.ArgumentCaptor;
42 import org.mockito.Mock;
43 import org.mockito.Mockito;
44 import org.mockito.MockitoAnnotations;
45
46 import java.io.ByteArrayInputStream;
47 import java.io.InputStream;
48 import java.util.List;
49
50 public class TestResponseEntityProxy {
51
52 @Mock
53 private ClassicHttpResponse response;
54 @Mock
55 private ExecRuntime execRuntime;
56 @Mock
57 private HttpEntity entity;
58
59 @Before
60 public void setUp() {
61 MockitoAnnotations.initMocks(this);
62 Mockito.when(entity.isStreaming()).thenReturn(Boolean.TRUE);
63 Mockito.when(response.getEntity()).thenReturn(entity);
64 }
65
66 @Test
67 public void testGetTrailersWithNoChunkedInputStream() throws Exception {
68 final ByteArrayInputStream inputStream = new ByteArrayInputStream("Test payload".getBytes());
69 Mockito.when(entity.getContent()).thenReturn(inputStream);
70 final ArgumentCaptor<HttpEntity> httpEntityArgumentCaptor = ArgumentCaptor.forClass(HttpEntity.class);
71
72 ResponseEntityProxy.enhance(response, execRuntime);
73
74 Mockito.verify(response).setEntity(httpEntityArgumentCaptor.capture());
75 final HttpEntity wrappedEntity = httpEntityArgumentCaptor.getValue();
76
77 final InputStream is = wrappedEntity.getContent();
78 while (is.read() != -1) {}
79 final Supplier<List<? extends Header>> trailers = wrappedEntity.getTrailers();
80
81 Assert.assertTrue(trailers.get().isEmpty());
82 }
83
84 @Test
85 public void testGetTrailersWithChunkedInputStream() throws Exception {
86 final SessionInputBuffer sessionInputBuffer = new SessionInputBufferImpl(100);
87 final ByteArrayInputStream inputStream = new ByteArrayInputStream("0\r\nX-Test-Trailer-Header: test\r\n".getBytes());
88 final ChunkedInputStream chunkedInputStream = new ChunkedInputStream(sessionInputBuffer, inputStream);
89
90 Mockito.when(entity.getContent()).thenReturn(chunkedInputStream);
91 final ArgumentCaptor<HttpEntity> httpEntityArgumentCaptor = ArgumentCaptor.forClass(HttpEntity.class);
92
93 ResponseEntityProxy.enhance(response, execRuntime);
94
95 Mockito.verify(response).setEntity(httpEntityArgumentCaptor.capture());
96 final HttpEntity wrappedEntity = httpEntityArgumentCaptor.getValue();
97
98 final InputStream is = wrappedEntity.getContent();
99 while (is.read() != -1) {}
100 final Supplier<List<? extends Header>> trailers = wrappedEntity.getTrailers();
101 final List<? extends Header> headers = trailers.get();
102
103 Assert.assertEquals(1, headers.size());
104 final Header header = headers.get(0);
105 Assert.assertEquals("X-Test-Trailer-Header", header.getName());
106 Assert.assertEquals("test", header.getValue());
107 }
108 }