org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody Java Examples

@Test
public void whenNoCursorsThenLatestOffsetsAreUsed() throws IOException, InvalidCursorException {
    when(eventTypeCache.getEventType(TEST_EVENT_TYPE_NAME)).thenReturn(EVENT_TYPE);
    final List<PartitionStatistics> tps2 = ImmutableList.of(
            new KafkaPartitionStatistics(timeline, 0, 0, 87),
            new KafkaPartitionStatistics(timeline, 1, 0, 34));
    when(timelineService.getActiveTimeline(any(EventType.class))).thenReturn(timeline);
    when(topicRepositoryMock.loadTopicStatistics(eq(Collections.singletonList(timeline))))
            .thenReturn(tps2);

    final ArgumentCaptor<EventStreamConfig> configCaptor = ArgumentCaptor.forClass(EventStreamConfig.class);
    final EventStream eventStreamMock = mock(EventStream.class);
    when(eventStreamFactoryMock.createEventStream(any(), any(), configCaptor.capture(), any()))
            .thenReturn(eventStreamMock);

    final StreamingResponseBody responseBody = createStreamingResponseBody(1, 0, 1, 1, 0, null);
    responseBody.writeTo(new ByteArrayOutputStream());

    final EventStreamConfig streamConfig = configCaptor.getValue();
    assertThat(
            streamConfig.getCursors(),
            equalTo(tps2.stream().map(PartitionStatistics::getLast).collect(Collectors.toList())));
}
 

/**
 * Add the following preparation in cache.
 *
 * @param preparation the preparation to cache.
 * @param stepId the preparation step id.
 */
private void addPreparationInCache(PreparationDTO preparation, String stepId) {
    final ExportParameters exportParameters = new ExportParameters();
    exportParameters.setPreparationId(preparation.getId());
    exportParameters.setExportType("JSON");
    exportParameters.setStepId(stepId);
    exportParameters.setDatasetId(preparation.getDataSetId());

    final StreamingResponseBody streamingResponseBody = executeSampleExportStrategy(exportParameters);
    try {
        // the result is not important here as it will be cached !
        streamingResponseBody.writeTo(new NullOutputStream());
    } catch (IOException e) {
        throw new TDPException(UNEXPECTED_EXCEPTION, e);
    }
}
 

@Override
public StreamingResponseBody execute(ExportParameters parameters) {
    final String formatName = parameters.getExportType();
    final ExportFormat format = getFormat(formatName);
    ExportUtils.setExportHeaders(parameters.getExportName(), //
            parameters.getArguments().get(ExportFormat.PREFIX + CSVFormat.ParametersCSV.ENCODING), //
            format);
    return outputStream -> {
        // get the dataset content (in an auto-closable block to make sure it is properly closed)
        final String datasetId = parameters.getDatasetId();
        try (DataSet dataSet = datasetClient.getDataSet(datasetId, false, true)) {
            // get the actions to apply (no preparation ==> dataset export ==> no actions)
            Configuration configuration = Configuration
                    .builder() //
                    .args(parameters.getArguments()) //
                    .outFilter(rm -> filterService.build(parameters.getFilter(), rm)) //
                    .format(format.getName()) //
                    .volume(Configuration.Volume.SMALL) //
                    .output(outputStream) //
                    .limit(limit) //
                    .build();
            factory.get(configuration).buildExecutable(dataSet, configuration).execute();
        }
    };
}
 

/**
 * Apply the preparation to the dataset out of the given IDs.
 *
 * @param preparationId the preparation id to apply on the dataset.
 * @param datasetId the dataset id to transform.
 * @param formatName The output {@link ExportFormat format}. This format also set the MIME response type.
 * @param stepId the preparation step id to use (default is 'head').
 * @param name the transformation name.
 * @param exportParams additional (optional) export parameters.
 */
//@formatter:off
@RequestMapping(value = "/apply/preparation/{preparationId}/dataset/{datasetId}/{format}", method = GET)
@ApiOperation(value = "Transform the given preparation to the given format on the given dataset id", notes = "This operation transforms the dataset using preparation id in the provided format.")
@VolumeMetered
public StreamingResponseBody applyOnDataset(@ApiParam(value = "Preparation id to apply.") @PathVariable(value = "preparationId") final String preparationId,
                                            @ApiParam(value = "DataSet id to transform.") @PathVariable(value = "datasetId") final String datasetId,
                                            @ApiParam(value = "Output format") @PathVariable("format") final String formatName,
                                            @ApiParam(value = "Step id", defaultValue = "head") @RequestParam(value = "stepId", required = false, defaultValue = "head") final String stepId,
                                            @ApiParam(value = "Name of the transformation", defaultValue = "untitled") @RequestParam(value = "name", required = false, defaultValue = "untitled") final String name,
                                            @RequestParam final Map<String, String> exportParams) {
    //@formatter:on
    final ExportParameters exportParameters = new ExportParameters();
    exportParameters.setPreparationId(preparationId);
    exportParameters.setDatasetId(datasetId);
    exportParameters.setExportType(formatName);
    exportParameters.setStepId(stepId);
    exportParameters.setExportName(name);
    exportParameters.getArguments().putAll(exportParams);

    return executeSampleExportStrategy(exportParameters);
}
 

@Test
void testDownload() {
  String fileId = "MyFileId";
  String contentType = "application/octet-stream";
  String filename = "filename";
  FileMeta fileMeta = mock(FileMeta.class);
  when(fileMeta.getContentType()).thenReturn(contentType);
  when(fileMeta.getFilename()).thenReturn(filename);
  when(dataService.findOneById("sys_FileMeta", fileId, FileMeta.class)).thenReturn(fileMeta);

  ResponseEntity<StreamingResponseBody> responseEntity = filesApiServiceImpl.download(fileId);
  assertEquals(OK, responseEntity.getStatusCode());
  assertEquals(valueOf(contentType), responseEntity.getHeaders().getContentType());
  assertEquals(
      parse("attachment; filename=\"filename\""),
      responseEntity.getHeaders().getContentDisposition());
}
 

@RequestMapping(value = "/apply", method = POST)
@ApiOperation(value = "Run the transformation given the provided export parameters",
        notes = "This operation transforms the dataset or preparation using parameters in export parameters.")
@VolumeMetered
@AsyncOperation(conditionalClass = GetPrepContentAsyncCondition.class, //
        resultUrlGenerator = PreparationGetContentUrlGenerator.class, //
        executionIdGeneratorClass = ExportParametersExecutionIdGenerator.class //
)
public StreamingResponseBody
        execute(@ApiParam(
                value = "Preparation id to apply.") @RequestBody @Valid @AsyncParameter @AsyncExecutionId final ExportParameters parameters)
                throws IOException {

    // Async behavior
    final ConditionalTest conditionalTest = applicationContext.getBean(GetPrepContentAsyncCondition.class);
    if (conditionalTest.apply(parameters)) {
        // write to cache
        executeSampleExportStrategy(parameters).writeTo(new NullOutputStream());
        return outputStream -> {
        };
    } else {
        // sync behavior
        return executeSampleExportStrategy(parameters);
    }
}
 

@Test
public void shouldUsedVersionedPreparation() throws IOException {
    // Given
    final ExportParameters parameters = new ExportParameters();
    parameters.setExportType("JSON");
    parameters.setPreparationId("prep-1234");
    parameters.setStepId("step-1234");

    final PreparationDTO preparation = new PreparationDTO();
    preparation.setId("prep-1234");
    preparation.setHeadId("step-1234");
    configurePreparation(preparation, "prep-1234", "step-1234");

    // When
    final StreamingResponseBody body = strategy.execute(parameters);
    body.writeTo(new NullOutputStream());

    // Then
    final ArgumentCaptor<Configuration> captor = ArgumentCaptor.forClass(Configuration.class);
    verify(transformer).buildExecutable(any(), captor.capture());
    assertEquals("prep-1234", captor.getValue().getPreparationId());
    assertEquals("step-1234", captor.getValue().getPreparation().getHeadId());
}
 

@Test
public void shouldUsedHeadPreparation() throws IOException {
    // Given
    final ExportParameters parameters = new ExportParameters();
    parameters.setExportType("JSON");
    parameters.setPreparationId("prep-1234");
    parameters.setStepId("head");

    final PreparationDTO preparation = new PreparationDTO();
    preparation.getSteps().add(Step.ROOT_STEP.id());
    preparation.setId("prep-1234");
    preparation.setHeadId("head");
    configurePreparation(preparation, "prep-1234", "head");

    // When
    final StreamingResponseBody body = strategy.execute(parameters);
    body.writeTo(new NullOutputStream());

    // Then
    final ArgumentCaptor<Configuration> captor = ArgumentCaptor.forClass(Configuration.class);
    verify(transformer).buildExecutable(any(), captor.capture());
    assertEquals("prep-1234", captor.getValue().getPreparationId());
    assertEquals("head", captor.getValue().getPreparation().getHeadId());
}
 

/**
 * Get the suggested action dynamic params. Dynamic params depends on the context (dataset / preparation / actual
 * transformations)
 */
@RequestMapping(value = "/api/transform/suggest/{action}/params", method = GET, produces = APPLICATION_JSON_VALUE)
@ApiOperation(value = "Get the transformation dynamic parameters", notes = "Returns the transformation parameters.")
@Timed
public ResponseEntity<StreamingResponseBody> suggestActionParams(
        @ApiParam(value = "Transformation name.") @PathVariable("action") final String action, @ApiParam(
                value = "Suggested dynamic transformation input (preparation id or dataset id") @Valid final DynamicParamsInput dynamicParamsInput) {
    // get preparation/dataset content
    HystrixCommand<InputStream> inputData;
    final String preparationId = dynamicParamsInput.getPreparationId();
    if (isNotBlank(preparationId)) {
        inputData = new AsyncGet<>(
                () -> getCommand(PreparationGetContent.class, preparationId, dynamicParamsInput.getStepId()),
                commonAPI);
    } else {
        inputData = datasetClient.getDataSetGetCommand(dynamicParamsInput.getDatasetId(), false, false);
    }

    // get params, passing content in the body
    final GenericCommand<InputStream> getActionDynamicParams =
            getCommand(SuggestActionParams.class, inputData, action, dynamicParamsInput.getColumnId());
    return CommandHelper.toStreaming(getActionDynamicParams);
}
 

@RequestMapping(value = "/api/preparations/preview/add", method = POST, consumes = APPLICATION_JSON_VALUE, produces = APPLICATION_JSON_VALUE)
@ApiOperation(value = "Get a preview between the head step and a new appended transformation")
public StreamingResponseBody previewAdd(@RequestBody @Valid final PreviewAddParameters input) {
//@formatter:on

    PreparationDTO preparation = null;
    List<Action> actions = new ArrayList<>(0);

    // get preparation details with dealing with preparations
    if (StringUtils.isNotBlank(input.getPreparationId())) {
        preparation = internalGetPreparation(input.getPreparationId());
        actions = internalGetActions(preparation.getId());
    }

    final HystrixCommand<InputStream> transformation = getCommand(PreviewAdd.class, input, preparation, actions);
    return executePreviewCommand(transformation);
}
 

@RequestMapping(value = "/api/datasets/preview/{id}", method = GET, produces = APPLICATION_JSON_VALUE)
@ApiOperation(value = "Get a data set by id.", produces = APPLICATION_JSON_VALUE,
        notes = "Get a data set based on given id.")
@Timed
public ResponseEntity<StreamingResponseBody> preview(
        @ApiParam(value = "Id of the data set to get") @PathVariable(value = "id") String id,
        @RequestParam(defaultValue = "true") @ApiParam(name = "metadata",
                value = "Include metadata information in the response") boolean metadata,
        @RequestParam(defaultValue = "") @ApiParam(name = "sheetName",
                value = "Sheet name to preview") String sheetName) {
    if (LOG.isDebugEnabled()) {
        LOG.debug("Requesting dataset #{} (pool: {})...", id, getConnectionStats());
    }
    try {
        GenericCommand<InputStream> retrievalCommand = getCommand(DataSetPreview.class, id, metadata, sheetName);
        return toStreaming(retrievalCommand);
    } finally {
        if (LOG.isDebugEnabled()) {
            LOG.debug("Request dataset #{} (pool: {}) done.", id, getConnectionStats());
        }
    }
}
 

public static StreamingResponseBody toStreaming(final HystrixCommand<InputStream> command) {
    return outputStream -> {
        final Observable<InputStream> stream = command.toObservable();
        stream.toBlocking().subscribe(inputStream -> {
            try {
                IOUtils.copyLarge(inputStream, outputStream);
                outputStream.flush();
            } catch (IOException e) {
                try {
                    inputStream.close();
                } catch (IOException closingException) {
                    LOGGER.warn("could not close command result, a http connection may be leaked !",
                            closingException);
                }
                LOGGER.error("Unable to fully copy command result '{}'.", command.getClass(), e);
            }
        }, TDPException::rethrowOrWrap);
    };
}
 

@RequestMapping(value = "/subscriptions/{subscription_id}/events", method = RequestMethod.GET)
public StreamingResponseBody streamEvents(
        @PathVariable("subscription_id") final String subscriptionId,
        @Nullable @RequestParam(value = "max_uncommitted_events", required = false) final Integer
                maxUncommittedEvents,
        @Nullable @RequestParam(value = "batch_limit", required = false) final Integer batchLimit,
        @Nullable @RequestParam(value = "stream_limit", required = false) final Long streamLimit,
        @Nullable @RequestParam(value = "batch_timespan", required = false) final Long batchTimespan,
        @Nullable @RequestParam(value = "batch_flush_timeout", required = false) final Integer batchTimeout,
        @Nullable @RequestParam(value = "stream_timeout", required = false) final Long streamTimeout,
        @Nullable @RequestParam(value = "stream_keep_alive_limit", required = false) final Integer
                streamKeepAliveLimit,
        @Nullable @RequestParam(value = "commit_timeout", required = false) final Long commitTimeout,
        final HttpServletRequest request, final HttpServletResponse response, final Client client) {
    final UserStreamParameters userParameters = new UserStreamParameters(batchLimit, streamLimit, batchTimespan,
            batchTimeout, streamTimeout, streamKeepAliveLimit, maxUncommittedEvents, ImmutableList.of(),
            commitTimeout);

    final StreamParameters streamParameters = StreamParameters.of(userParameters,
            nakadiSettings.getMaxCommitTimeout(), client);

    return stream(subscriptionId, request, response, client, streamParameters,
            TracingService.extractSpan(request, "stream_events_request")
                    .setTag("subscription.id", subscriptionId));
}
 

/**
 * Get the current dictionary (as serialized object).
 */
@RequestMapping(value = "/api/transform/dictionary", method = GET, produces = APPLICATION_JSON_VALUE)
@ApiOperation(value = "Get current dictionary (as serialized object).",
        notes = "Returns a DQ dictionary serialized usin Java serialization and GZIP-ed.")
@Timed
public StreamingResponseBody getDictionary() {
    // get preparation/dataset content
    HystrixCommand<InputStream> dictionaryCommand = getCommand(DictionaryCommand.class);
    return CommandHelper.toStreaming(dictionaryCommand);
}
 

@RequestMapping(value = "/api/folders/{id}", method = GET, produces = APPLICATION_JSON_VALUE)
@ApiOperation(value = "Get folder by id", produces = APPLICATION_JSON_VALUE, notes = "Get a folder by id")
@Timed
public ResponseEntity<StreamingResponseBody>
        getFolderAndHierarchyById(@PathVariable(value = "id") final String id) {
    try {
        final HystrixCommand<InputStream> foldersList = getCommand(GetFolder.class, id);
        return ResponseEntity
                .ok() //
                .contentType(APPLICATION_JSON_UTF8) //
                .body(CommandHelper.toStreaming(foldersList));
    } catch (Exception e) {
        throw new TDPException(APIErrorCodes.UNABLE_TO_GET_FOLDERS, e);
    }
}
 

/**
 * Return the semantic types for a given preparation / column.
 *
 * @param preparationId the preparation id.
 * @param columnId the column id.
 * @param stepId the step id (optional, if not specified, it's 'head')
 * @return the semantic types for a given preparation / column.
 */
@RequestMapping(value = "/api/preparations/{preparationId}/columns/{columnId}/types", method = GET,
        produces = APPLICATION_JSON_VALUE)
@ApiOperation(value = "list the types of the wanted column",
        notes = "This list can be used by user to change the column type.")
@Timed
@PublicAPI
public ResponseEntity<StreamingResponseBody> getPreparationColumnSemanticCategories(
        @ApiParam(value = "The preparation id") @PathVariable String preparationId,
        @ApiParam(value = "The column id") @PathVariable String columnId,
        @ApiParam(value = "The preparation version") @RequestParam(defaultValue = "head") String stepId) {

    LOG.debug("listing semantic types for preparation {} / {}, column {}", preparationId, columnId, stepId);
    return CommandHelper.toStreaming(getCommand(GetPreparationColumnTypes.class, preparationId, columnId, stepId));
}
 

/**
 * no javadoc here so see description in @ApiOperation notes.
 *
 * @param name The folder to search.
 * @param strict Strict mode means searched name is the full name.
 * @return the list of folders that match the given name.
 */
@RequestMapping(value = "/api/folders/search", method = GET, produces = APPLICATION_JSON_VALUE)
@ApiOperation(value = "Search Folders with parameter as part of the name", produces = APPLICATION_JSON_VALUE)
@Timed
public ResponseEntity<StreamingResponseBody> search(@RequestParam(required = false) final String name,
        @RequestParam(required = false) final Boolean strict, @RequestParam(required = false) final String path) {
    try {
        final GenericCommand<InputStream> searchFolders = getCommand(SearchFolders.class, name, strict, path);
        return CommandHelper.toStreaming(searchFolders);
    } catch (Exception e) {
        throw new TDPException(APIErrorCodes.UNABLE_TO_LIST_FOLDERS, e);
    }
}
 

@RequestMapping(value = "/api/folders", method = PUT)
@ApiOperation(value = "Add a folder.", produces = APPLICATION_JSON_VALUE)
@Timed
public StreamingResponseBody addFolder(@RequestParam(required = false) final String parentId,
        @RequestParam final String path) {
    try {
        final HystrixCommand<InputStream> createChildFolder = getCommand(CreateChildFolder.class, parentId, path);
        return CommandHelper.toStreaming(createChildFolder);
    } catch (Exception e) {
        throw new TDPException(APIErrorCodes.UNABLE_TO_CREATE_FOLDER, e);
    }
}
 

@RequestMapping(value = "/api/folders", method = GET)
@ApiOperation(value = "List folders. Optional filter on parent ID may be supplied.",
        produces = APPLICATION_JSON_VALUE)
@Timed
public ResponseEntity<StreamingResponseBody> listFolders(@RequestParam(required = false) String parentId) {
    try {
        final GenericCommand<InputStream> foldersList = getCommand(FolderChildrenList.class, parentId);
        return CommandHelper.toStreaming(foldersList);
    } catch (Exception e) {
        throw new TDPException(APIErrorCodes.UNABLE_TO_LIST_FOLDERS, e);
    }
}
 

/**
 * Compute an aggregation according to the given parameters.
 *
 * @param input The aggregation parameters.
 */
@RequestMapping(value = "/api/aggregate", method = POST, consumes = APPLICATION_JSON_VALUE,
        produces = APPLICATION_JSON_VALUE)
@ApiOperation(value = "Compute aggregation", consumes = APPLICATION_JSON_VALUE, produces = APPLICATION_JSON_VALUE,
        notes = "Compute aggregation according to the given parameters")
public ResponseEntity<StreamingResponseBody> compute(@RequestBody @Valid final AggregationParameters input) {
    LOG.debug("Aggregation computation requested (pool: {} )...", getConnectionStats());
    // get the command and execute it, then copy the content to the http response
    try {
        GenericCommand<InputStream> command = getCommand(Aggregate.class, input);
        return CommandHelper.toStreaming(command);
    } finally {
        LOG.debug("Aggregation done (pool: {} )...", getConnectionStats());
    }
}
 

@ApiOperation("Download file (see documentation)")
@GetMapping(value = "/{fileId}", params = "alt=media")
public ResponseEntity<StreamingResponseBody> downloadFile(@PathVariable("fileId") String fileId) {
  validateReadPermission();

  return filesService.download(fileId);
}
 

@RequestMapping(value = "/api/folders/tree", method = GET)
@ApiOperation(value = "List all folders", produces = APPLICATION_JSON_VALUE)
@Timed
public StreamingResponseBody getTree() {
    try {
        final HystrixCommand<InputStream> foldersList = getCommand(FolderTree.class);
        return CommandHelper.toStreaming(foldersList);
    } catch (Exception e) {
        throw new TDPException(APIErrorCodes.UNABLE_TO_LIST_FOLDERS, e);
    }
}
 

@Override
public StreamingResponseBody execute(ExportParameters parameters) {
    final TransformationCacheKey contentKey = getCacheKey(parameters);
    ExportUtils.setExportHeaders(parameters.getExportName(), //
            parameters.getArguments().get(ExportFormat.PREFIX + CSVFormat.ParametersCSV.ENCODING), //
            getFormat(parameters.getExportType()));
    LOGGER.debug("Using '{}' content cache entry.", contentKey.getKey());
    return outputStream -> {
        try (InputStream cachedContent = contentCache.get(contentKey)) {
            IOUtils.copy(cachedContent, outputStream);
        }
    };
}
 

@Test
public void whenNakadiExceptionIsThrownThenServiceUnavailable() throws IOException {
    when(eventTypeCache.getEventType(TEST_EVENT_TYPE_NAME))
            .thenThrow(ServiceTemporarilyUnavailableException.class);

    final StreamingResponseBody responseBody = createStreamingResponseBody();

    final Problem expectedProblem = Problem.valueOf(SERVICE_UNAVAILABLE);
    MatcherAssert.assertThat(responseToString(responseBody), JSON_TEST_HELPER.matchesObject(expectedProblem));
}
 

@GetMapping(path = "/sitemap.xml", produces = MediaType.APPLICATION_XML_VALUE)
public ResponseEntity<StreamingResponseBody> getSitemap() throws ParserConfigurationException {
    var document = DocumentBuilderFactory.newInstance().newDocumentBuilder().newDocument();
    document.setXmlStandalone(true);
    var urlset = document.createElementNS(NAMESPACE_URI, "urlset");
    document.appendChild(urlset);
    
    var baseUrl = getBaseUrl();
    var timestampFormatter = DateTimeFormatter.ofPattern("yyyy-MM-dd");
    repositories.findAllExtensions().forEach(extension -> {
        var entry = document.createElement("url");
        var loc = document.createElement("loc");
        loc.setTextContent(UrlUtil.createApiUrl(baseUrl, "extension", extension.getNamespace().getName(), extension.getName()));
        entry.appendChild(loc);
        var lastmod = document.createElement("lastmod");
        lastmod.setTextContent(extension.getLatest().getTimestamp().format(timestampFormatter));
        entry.appendChild(lastmod);
        urlset.appendChild(entry);
    });

    StreamingResponseBody stream = out -> {
        try {
            var transformer = TransformerFactory.newInstance().newTransformer();
            transformer.setOutputProperty(OutputKeys.ENCODING, "UTF-8");
            transformer.setOutputProperty(OutputKeys.INDENT, "yes");
            transformer.setOutputProperty("{http://xml.apache.org/xslt}indent-amount", "2");
            transformer.transform(new DOMSource(document), new StreamResult(out));
        } catch (TransformerException exc) {
            throw new RuntimeException(exc);
        }
    };
    return new ResponseEntity<>(stream, HttpStatus.OK);
}
 

@Test
public void testAccessDenied() throws Exception {
    Mockito.doThrow(AccessDeniedException.class).when(authorizationValidator)
            .authorizeStreamRead(any());

    when(eventTypeCache.getEventType(TEST_EVENT_TYPE_NAME)).thenReturn(EVENT_TYPE);
    Mockito.doThrow(mockAccessDeniedException()).when(authorizationValidator).authorizeStreamRead(any());

    final StreamingResponseBody responseBody = createStreamingResponseBody(0, 0, 0, 0, 0, null);

    final Problem expectedProblem = Problem.valueOf(FORBIDDEN, "Access on READ some-type:some-name denied");
    MatcherAssert.assertThat(responseToString(responseBody), JSON_TEST_HELPER.matchesObject(expectedProblem));
}
 

private StreamingResponseBody createStreamingResponseBody(final Integer batchLimit,
                                                          final Integer streamLimit,
                                                          final Integer batchTimeout,
                                                          final Integer streamTimeout,
                                                          final Integer streamKeepAliveLimit,
                                                          final String cursorsStr) throws IOException {
    return controller.streamEvents(TEST_EVENT_TYPE_NAME, batchLimit, streamLimit, batchTimeout, streamTimeout,
            streamKeepAliveLimit, cursorsStr, requestMock, responseMock, FULL_ACCESS_CLIENT);
}
 

public static ResponseEntity<StreamingResponseBody> toStreaming(final GenericCommand<InputStream> command) {
    final Observable<InputStream> stream = command.toObservable();
    return stream.map(is -> {
        // Content for the response entity
        final StreamingResponseBody body = outputStream -> {
            try {
                IOUtils.copyLarge(is, outputStream);
                outputStream.flush();
            } catch (IOException e) {
                try {
                    is.close();
                } catch (IOException closingException) {
                    LOGGER.warn("could not close command result, a http connection may be leaked !",
                            closingException);
                }
                LOGGER.error("Unable to fully copy command result '{}'.", command.getClass(), e);
            }
        };
        // copy all headers from the command response so that the mime-type is correctly forwarded. Command has
        // the correct headers due to call to toBlocking() below.
        final MultiValueMap<String, String> headers = new HttpHeaders();
        final HttpStatus status = command.getStatus();
        for (Header header : command.getCommandResponseHeaders()) {
            headers.put(header.getName(), Collections.singletonList(header.getValue()));
        }
        return new ResponseEntity<>(body, headers, status == null ? HttpStatus.OK : status);
    }).toBlocking().first();
}
 

@Test
public void shouldUpdateStepRowMetadataForEachPersistentPreparation() throws Exception {
    // given
    when(folderRepository.locateEntry(any(), any())).thenReturn(mock(Folder.class));
    when(service.execute(any(ExportParameters.class))).thenReturn(mock(StreamingResponseBody.class));

    // when
    task.run();

    // then
    verify(service, times(preparations.size())).execute(any());
}
 

@Test
public void testCommandToStreamingWithHeader() throws Exception {
    GenericCommand<InputStream> command = new CommandHelperTestCommand();

    final ResponseEntity<StreamingResponseBody> responseEntity = CommandHelper.toStreaming(command);
    final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
    responseEntity.getBody().writeTo(outputStream);
    assertEquals("test", new String(outputStream.toByteArray()));
    assertEquals(HttpStatus.NO_CONTENT, responseEntity.getStatusCode());
    assertEquals("custom value", responseEntity.getHeaders().get("custom").get(0));
}