Rest proxy remove block from Schedulers (#5376)

* Fixed the case when deserialized header is null
Azure · Sep 16, 2019 · 4bb9450 · 4bb9450
1 parent b8b555d
commit 4bb9450
Showing 1 changed file with 34 additions and 23 deletions.
diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/implementation/RestProxy.java b/sdk/core/azure-core/src/main/java/com/azure/core/implementation/RestProxy.java
@@ -459,13 +459,13 @@ private Mono<?> handleRestResponseReturnType(HttpDecodedResponse response, Swagg
 
             if (TypeUtil.isTypeOrSubTypeOf(bodyType, Void.class)) {
                 asyncResult = response.getSourceResponse().getBody().ignoreElements()
-                    .then(Mono.just(createResponse(response, entityType, null)));
+                    .then(createResponse(response, entityType, null));
             } else {
                 asyncResult = handleBodyReturnType(response, methodParser, bodyType)
-                    .map((Function<Object, Response<?>>) bodyAsObject -> createResponse(response, entityType,
+                    .flatMap((Function<Object, Mono<Response<?>>>) bodyAsObject -> createResponse(response, entityType,
                         bodyAsObject))
-                    .switchIfEmpty(Mono.defer((Supplier<Mono<Response<?>>>) () -> Mono.just(createResponse(response,
-                        entityType, null))));
+                    .switchIfEmpty(Mono.defer((Supplier<Mono<Response<?>>>) () -> createResponse(response,
+                        entityType, null)));
             }
         } else {
             // For now we're just throwing if the Maybe didn't emit a value.
@@ -476,7 +476,7 @@ private Mono<?> handleRestResponseReturnType(HttpDecodedResponse response, Swagg
     }
 
     @SuppressWarnings("unchecked")
-    private Response<?> createResponse(HttpDecodedResponse response, Type entityType, Object bodyAsObject) {
+    private Mono<Response<?>> createResponse(HttpDecodedResponse response, Type entityType, Object bodyAsObject) {
         final HttpResponse httpResponse = response.getSourceResponse();
         final HttpRequest httpRequest = httpResponse.getRequest();
         final int responseStatusCode = httpResponse.getStatusCode();
@@ -515,30 +515,41 @@ private Response<?> createResponse(HttpDecodedResponse response, Type entityType
         // try to create an instance using our list of potential candidates
         for (Constructor<?> constructor : constructors) {
             final Constructor<? extends Response<?>> ctor = (Constructor<? extends Response<?>>) constructor;
-
-            try {
-                final int paramCount = constructor.getParameterCount();
-
-                switch (paramCount) {
-                    case 3:
-                        return ctor.newInstance(httpRequest, responseStatusCode, responseHeaders);
-                    case 4:
-                        return ctor.newInstance(httpRequest, responseStatusCode, responseHeaders, bodyAsObject);
-                    case 5:
-                        return ctor.newInstance(httpRequest, responseStatusCode, responseHeaders, bodyAsObject,
-                            response.getDecodedHeaders().block());
-                    default:
-                        throw logger.logExceptionAsError(new IllegalStateException(
-                            "Response constructor with expected parameters not found."));
-                }
-            } catch (IllegalAccessException | InvocationTargetException | InstantiationException e) {
-                throw logger.logExceptionAsError(reactor.core.Exceptions.propagate(e));
+            final int paramCount = constructor.getParameterCount();
+
+            switch (paramCount) {
+                case 3:
+                    return Mono.just(createResponse(ctor, new Object[]
+                        {httpRequest, responseStatusCode, responseHeaders}));
+                case 4:
+                    return Mono.just(createResponse(ctor, new Object[] {httpRequest, responseStatusCode,
+                        responseHeaders, bodyAsObject}));
+                case 5:
+                    return response.getDecodedHeaders()
+                        .map((Function<Object, Response<?>>) headers -> {
+                            return createResponse(ctor, new Object[]
+                                {httpRequest, responseStatusCode, responseHeaders, bodyAsObject, headers});
+                        }).switchIfEmpty(Mono.defer((Supplier<Mono<Response<?>>>) () -> {
+                            return Mono.just(createResponse(ctor, new Object[]
+                                {httpRequest, responseStatusCode, responseHeaders, bodyAsObject, null}));
+                        }));
+                default:
+                    throw logger.logExceptionAsError(new IllegalStateException(
+                        "Response constructor with expected parameters not found."));
             }
         }
         // error
         throw logger.logExceptionAsError(new RuntimeException("Cannot find suitable constructor for class " + cls));
     }
 
+    private Response<?> createResponse(Constructor<? extends Response<?>> ctor, Object[] args) {
+        try {
+            return ctor.newInstance(args);
+        } catch (IllegalAccessException | InvocationTargetException | InstantiationException e) {
+            throw logger.logExceptionAsError(Exceptions.propagate(e));
+        }
+    }
+
     protected final Mono<?> handleBodyReturnType(final HttpDecodedResponse response,
                 final SwaggerMethodParser methodParser, final Type entityType) {
         final int responseStatusCode = response.getSourceResponse().getStatusCode();