using ResponseEntity in ProcessInstanceDiagramResource

modules/activiti-rest/src/main/java/org/activiti/rest/service/api/runtime/process/ProcessInstanceDiagramResource.java

@@ -22,16 +22,17 @@ import org.activiti.bpmn.model.BpmnModel;
 import org.activiti.engine.ActivitiIllegalArgumentException;
 import org.activiti.engine.ProcessEngineConfiguration;
 import org.activiti.engine.RepositoryService;
-import org.activiti.engine.impl.RepositoryServiceImpl;
-import org.activiti.engine.impl.persistence.entity.ProcessDefinitionEntity;
+import org.activiti.engine.repository.ProcessDefinition;
 import org.activiti.engine.runtime.ProcessInstance;
 import org.activiti.image.ProcessDiagramGenerator;
 import org.apache.commons.io.IOUtils;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.PathVariable;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RequestMethod;
-import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
 /**
@@ -47,31 +48,27 @@ public class ProcessInstanceDiagramResource extends BaseProcessInstanceResource
   protected ProcessEngineConfiguration processEngineConfiguration;
   
   @RequestMapping(value="/runtime/process-instances/{processInstanceId}/diagram", method = RequestMethod.GET)
-  public @ResponseBody byte[] getProcessInstanceDiagram(@PathVariable String processInstanceId, HttpServletResponse response) {
+  public ResponseEntity<byte[]> getProcessInstanceDiagram(@PathVariable String processInstanceId, HttpServletResponse response) {
     ProcessInstance processInstance = getProcessInstanceFromRequest(processInstanceId);
     
-    ProcessDefinitionEntity pde = (ProcessDefinitionEntity) ((RepositoryServiceImpl) repositoryService)
-        .getDeployedProcessDefinition(processInstance.getProcessDefinitionId());
+    ProcessDefinition pde = repositoryService.getProcessDefinition(processInstance.getProcessDefinitionId());
 
-    if (pde != null && pde.isGraphicalNotationDefined()) {
+    if (pde != null && pde.hasGraphicalNotation()) {
       BpmnModel bpmnModel = repositoryService.getBpmnModel(pde.getId());
       ProcessDiagramGenerator diagramGenerator = processEngineConfiguration.getProcessDiagramGenerator();
       InputStream resource = diagramGenerator.generateDiagram(bpmnModel, "png", runtimeService.getActiveActivityIds(processInstance.getId()),
           Collections.<String>emptyList(), processEngineConfiguration.getActivityFontName(), processEngineConfiguration.getLabelFontName(),
           processEngineConfiguration.getClassLoader(), 1.0);
 
+      HttpHeaders responseHeaders = new HttpHeaders();
+      responseHeaders.set("Content-Type", "image/png");
       try {
-        byte[] responseBytes = IOUtils.toByteArray(resource);
-        response.setContentType("image/png");
-        return responseBytes;
-        
+        return new ResponseEntity<byte[]>(IOUtils.toByteArray(resource), responseHeaders, HttpStatus.OK);
       } catch(Exception e) {
-        response.setContentType("application/json");
         throw new ActivitiIllegalArgumentException("Error exporting diagram", e);
       }
       
     } else {
-      response.setContentType("application/json");
       throw new ActivitiIllegalArgumentException("Process instance with id '" + processInstance.getId() + "' has no graphical notation defined.");
     }
   }

modules/activiti-rest/src/test/java/org/activiti/rest/service/api/runtime/ProcessInstanceDiagramResourceTest.java

@@ -34,6 +34,7 @@ public class ProcessInstanceDiagramResourceTest extends BaseSpringRestTestCase {
     CloseableHttpResponse response = executeRequest(new HttpGet(SERVER_URL_PREFIX + 
         RestUrls.createRelativeResourceUrl(RestUrls.URL_PROCESS_INSTANCE_DIAGRAM, processInstance.getId())), HttpStatus.SC_OK);
     assertNotNull(response.getEntity().getContent());
+    assertEquals("image/png", response.getEntity().getContentType().getValue());
     closeResponse(response);
   }