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 java.io.ByteArrayInputStream;
31 import java.io.InputStream;
32 import java.util.List;
33
34 import org.apache.hc.client5.http.classic.ExecRuntime;
35 import org.apache.hc.core5.function.Supplier;
36 import org.apache.hc.core5.http.ClassicHttpResponse;
37 import org.apache.hc.core5.http.Header;
38 import org.apache.hc.core5.http.HttpEntity;
39 import org.apache.hc.core5.http.impl.io.ChunkedInputStream;
40 import org.apache.hc.core5.http.impl.io.SessionInputBufferImpl;
41 import org.apache.hc.core5.http.io.SessionInputBuffer;
42 import org.apache.hc.core5.http.io.entity.BasicHttpEntity;
43 import org.apache.hc.core5.http.message.BasicClassicHttpResponse;
44 import org.junit.jupiter.api.Assertions;
45 import org.junit.jupiter.api.BeforeEach;
46 import org.junit.jupiter.api.Test;
47 import org.mockito.ArgumentCaptor;
48 import org.mockito.Mock;
49 import org.mockito.Mockito;
50 import org.mockito.MockitoAnnotations;
51
52 public class TestResponseEntityProxy {
53
54 @Mock
55 private ClassicHttpResponse response;
56 @Mock
57 private ExecRuntime execRuntime;
58 @Mock
59 private HttpEntity entity;
60
61 @BeforeEach
62 public void setUp() {
63 MockitoAnnotations.openMocks(this);
64 Mockito.when(entity.isStreaming()).thenReturn(Boolean.TRUE);
65 Mockito.when(response.getEntity()).thenReturn(entity);
66 }
67
68 @Test
69 public void testGetTrailersWithNoChunkedInputStream() throws Exception {
70 final ByteArrayInputStream inputStream = new ByteArrayInputStream("Test payload".getBytes());
71 Mockito.when(entity.getContent()).thenReturn(inputStream);
72 final ArgumentCaptor<HttpEntity> httpEntityArgumentCaptor = ArgumentCaptor.forClass(HttpEntity.class);
73
74 ResponseEntityProxy.enhance(response, execRuntime);
75
76 Mockito.verify(response).setEntity(httpEntityArgumentCaptor.capture());
77 final HttpEntity wrappedEntity = httpEntityArgumentCaptor.getValue();
78
79 final InputStream is = wrappedEntity.getContent();
80 while (is.read() != -1) {}
81 final Supplier<List<? extends Header>> trailers = wrappedEntity.getTrailers();
82
83 Assertions.assertTrue(trailers.get().isEmpty());
84 }
85
86 @Test
87 public void testGetTrailersWithChunkedInputStream() throws Exception {
88 final SessionInputBuffer sessionInputBuffer = new SessionInputBufferImpl(100);
89 final ByteArrayInputStream inputStream = new ByteArrayInputStream("0\r\nX-Test-Trailer-Header: test\r\n".getBytes());
90 final ChunkedInputStream chunkedInputStream = new ChunkedInputStream(sessionInputBuffer, inputStream);
91
92 Mockito.when(entity.getContent()).thenReturn(chunkedInputStream);
93 final ArgumentCaptor<HttpEntity> httpEntityArgumentCaptor = ArgumentCaptor.forClass(HttpEntity.class);
94
95 ResponseEntityProxy.enhance(response, execRuntime);
96
97 Mockito.verify(response).setEntity(httpEntityArgumentCaptor.capture());
98 final HttpEntity wrappedEntity = httpEntityArgumentCaptor.getValue();
99
100 final InputStream is = wrappedEntity.getContent();
101 while (is.read() != -1) {}
102 final Supplier<List<? extends Header>> trailers = wrappedEntity.getTrailers();
103 final List<? extends Header> headers = trailers.get();
104
105 Assertions.assertEquals(1, headers.size());
106 final Header header = headers.get(0);
107 Assertions.assertEquals("X-Test-Trailer-Header", header.getName());
108 Assertions.assertEquals("test", header.getValue());
109 }
110
111 @Test
112 public void testWriteToNullDrainsAndReleasesStream() throws Exception {
113 final SessionInputBuffer sessionInputBuffer = new SessionInputBufferImpl(100);
114 final ByteArrayInputStream inputStream = new ByteArrayInputStream("0\r\nX-Test-Trailer-Header: test\r\n".getBytes());
115 final ChunkedInputStream chunkedInputStream = new ChunkedInputStream(sessionInputBuffer, inputStream);
116 final CloseableHttpResponse resp = new CloseableHttpResponse(new BasicClassicHttpResponse(200), execRuntime);
117 final HttpEntity entity = new BasicHttpEntity(chunkedInputStream, null, true);
118 Assertions.assertTrue(entity.isStreaming());
119 resp.setEntity(entity);
120
121 ResponseEntityProxy.enhance(resp, execRuntime);
122
123 final HttpEntity wrappedEntity = resp.getEntity();
124
125 wrappedEntity.writeTo(null);
126 Mockito.verify(execRuntime).releaseEndpoint();
127
128 final Supplier<List<? extends Header>> trailers = wrappedEntity.getTrailers();
129 final List<? extends Header> headers = trailers.get();
130
131 Assertions.assertEquals(1, headers.size());
132 final Header header = headers.get(0);
133 Assertions.assertEquals("X-Test-Trailer-Header", header.getName());
134 Assertions.assertEquals("test", header.getValue());
135
136
137 }
138 }