25
25
import javax .servlet .http .HttpServletRequest ;
26
26
import javax .servlet .http .HttpServletResponse ;
27
27
28
+ import org .springframework .mock .web .MockAsyncContext ;
28
29
import org .springframework .web .context .WebApplicationContext ;
29
30
import org .springframework .web .context .request .NativeWebRequest ;
30
- import org .springframework .web .context .request .async .CallableProcessingInterceptorAdapter ;
31
- import org .springframework .web .context .request .async .DeferredResult ;
32
- import org .springframework .web .context .request .async .DeferredResultProcessingInterceptorAdapter ;
33
- import org .springframework .web .context .request .async .WebAsyncManager ;
34
- import org .springframework .web .context .request .async .WebAsyncUtils ;
31
+ import org .springframework .web .context .request .async .*;
35
32
import org .springframework .web .servlet .DispatcherServlet ;
36
33
import org .springframework .web .servlet .HandlerExecutionChain ;
37
34
import org .springframework .web .servlet .ModelAndView ;
@@ -50,44 +47,52 @@ final class TestDispatcherServlet extends DispatcherServlet {
50
47
51
48
private static final String KEY = TestDispatcherServlet .class .getName () + ".interceptor" ;
52
49
50
+
53
51
/**
54
52
* Create a new instance with the given web application context.
55
53
*/
56
54
public TestDispatcherServlet (WebApplicationContext webApplicationContext ) {
57
55
super (webApplicationContext );
58
56
}
59
57
58
+
60
59
@ Override
61
60
protected void service (HttpServletRequest request , HttpServletResponse response ) throws ServletException , IOException {
62
61
63
- CountDownLatch latch = registerAsyncInterceptors (request );
64
- getMvcResult (request ).setAsyncResultLatch (latch );
62
+ registerAsyncResultInterceptors (request );
65
63
66
64
super .service (request , response );
67
- }
68
65
69
- private CountDownLatch registerAsyncInterceptors (final HttpServletRequest servletRequest ) {
70
-
71
- final CountDownLatch asyncResultLatch = new CountDownLatch (1 );
72
-
73
- WebAsyncManager asyncManager = WebAsyncUtils .getAsyncManager (servletRequest );
66
+ if (request .isAsyncStarted ()) {
67
+ addAsyncResultLatch (request );
68
+ }
69
+ }
74
70
71
+ private void registerAsyncResultInterceptors (final HttpServletRequest request ) {
72
+ WebAsyncManager asyncManager = WebAsyncUtils .getAsyncManager (request );
75
73
asyncManager .registerCallableInterceptor (KEY , new CallableProcessingInterceptorAdapter () {
76
74
@ Override
77
- public <T > void postProcess (NativeWebRequest request , Callable <T > task , Object value ) throws Exception {
78
- getMvcResult (servletRequest ).setAsyncResult (value );
79
- asyncResultLatch .countDown ();
75
+ public <T > void postProcess (NativeWebRequest r , Callable <T > task , Object value ) throws Exception {
76
+ getMvcResult (request ).setAsyncResult (value );
80
77
}
81
78
});
82
79
asyncManager .registerDeferredResultInterceptor (KEY , new DeferredResultProcessingInterceptorAdapter () {
83
80
@ Override
84
- public <T > void postProcess (NativeWebRequest request , DeferredResult <T > result , Object value ) throws Exception {
85
- getMvcResult (servletRequest ).setAsyncResult (value );
86
- asyncResultLatch .countDown ();
81
+ public <T > void postProcess (NativeWebRequest r , DeferredResult <T > result , Object value ) throws Exception {
82
+ getMvcResult (request ).setAsyncResult (value );
87
83
}
88
84
});
85
+ }
89
86
90
- return asyncResultLatch ;
87
+ private void addAsyncResultLatch (HttpServletRequest request ) {
88
+ final CountDownLatch latch = new CountDownLatch (1 );
89
+ ((MockAsyncContext ) request .getAsyncContext ()).addDispatchHandler (new Runnable () {
90
+ @ Override
91
+ public void run () {
92
+ latch .countDown ();
93
+ }
94
+ });
95
+ getMvcResult (request ).setAsyncResultLatch (latch );
91
96
}
92
97
93
98
protected DefaultMvcResult getMvcResult (ServletRequest request ) {
0 commit comments