MultipartResourceTest

package it.com.atlassian.rest.multipart;

import java.util.ArrayList;
import java.util.Iterator;

import javax.ws.rs.core.Response;
import javax.xml.bind.JAXBContext;
import javax.xml.bind.Unmarshaller;

import com.atlassian.plugins.rest.common.security.jersey.XsrfResourceFilter;
import com.atlassian.plugins.rest.multipart.FilePartObject;
import com.atlassian.plugins.rest.multipart.FilePartObjects;
import com.atlassian.rest.jersey.client.WebResourceFactory;

import org.apache.commons.httpclient.HttpClient;
import org.apache.commons.httpclient.methods.PostMethod;
import org.apache.commons.httpclient.methods.multipart.ByteArrayPartSource;
import org.apache.commons.httpclient.methods.multipart.FilePart;
import org.apache.commons.httpclient.methods.multipart.MultipartRequestEntity;
import org.apache.commons.httpclient.methods.multipart.Part;
import org.apache.commons.httpclient.methods.multipart.PartSource;
import org.apache.commons.httpclient.methods.multipart.StringPart;
import org.apache.commons.httpclient.params.HttpMethodParams;
import org.junit.Before;
import org.junit.Test;

import static junit.framework.Assert.assertTrue;
import static org.junit.Assert.assertEquals;


public class MultipartResourceTest {
    private HttpClient client;

    @Before
    public void setUp() {
        client = new HttpClient();
    }

    @Test
    public void testUploadSingleFile() throws Exception {
        FilePartObject filePartObject = new FilePartObject("file", false, "filename.txt", "text/plain", "Hello world!!");
        FilePartObject result = (FilePartObject) uploadAndUnmarshalFileParts("single", filePartObject);
        assertFilePartEquals(filePartObject, result);
    }

    @Test
    public void testUploadSingleSimpleField() throws Exception {
        FilePartObject filePartObject = new FilePartObject("file", true, null, null, "Hello world!!");
        FilePartObject result = (FilePartObject) uploadAndUnmarshalFileParts("single", filePartObject);
        assertFilePartEquals(filePartObject, result);
    }

    @Test
    public void testUploadMultiple() throws Exception {
        FilePartObject filePartObject1 = new FilePartObject("file", false, "file1.txt", "text/plain", "Hello first world!");
        FilePartObject filePartObject2 = new FilePartObject("file", false, "file2.txt", "text/plain", "Hello second world!");

        FilePartObjects result = (FilePartObjects) uploadAndUnmarshalFileParts("multiple", filePartObject1, filePartObject2);
        assertEquals(2, result.getFileParts().size());
        Iterator<FilePartObject> iter = result.getFileParts().iterator();
        assertFilePartEquals(filePartObject1, iter.next());
        assertFilePartEquals(filePartObject2, iter.next());
    }

    @Test
    public void testMaxFileSizeUnder() throws Exception {
        FilePartObject filePartObject = new FilePartObject("file", false, "file.txt", "text/plain", "Safe");
        FilePartObjects result = (FilePartObjects) uploadAndUnmarshalFileParts("config", filePartObject);
        assertEquals(1, result.getFileParts().size());
        assertFilePartEquals(filePartObject, result.getFileParts().iterator().next());
    }

    @Test
    public void testMaxFileSizeExceeded() throws Exception {
        // The config resource is configured for maximum file size of 10 characters
        FilePartObject filePartObject = new FilePartObject("file", false, "file.txt", "text/plain", "This is longer than 10 characters");
        PostMethod method = uploadFileParts("config", filePartObject);
        assertEquals(413, method.getStatusCode());
        assertTrue(method.getResponseBodyAsString().contains("10"));
    }

    @Test
    public void testMaxSizeExceeded() throws Exception {
        // The config resource is configured for maximum overall request size of 1000 characters.  This should yield
        // 1161 characters due to all the multipart headers
        PostMethod method = uploadFileParts("config", new FilePartObject("file", false, "file.txt", "text/plain", "Safe"),
                new FilePartObject("file", false, "file.txt", "text/plain", "Safe"),
                new FilePartObject("file", false, "file.txt", "text/plain", "Safe"),
                new FilePartObject("file", false, "file.txt", "text/plain", "Safe"),
                new FilePartObject("file", false, "file.txt", "text/plain", "Safe"),
                new FilePartObject("file", false, "file.txt", "text/plain", "Safe"));
        assertEquals(413, method.getStatusCode());
        assertTrue(method.getResponseBodyAsString().contains("1000"));
    }

    @Test
    public void testInterceptorInvoked() throws Exception {
        FilePartObject filePartObject = new FilePartObject("file", false, "abc.txt", "text/plain", "Upload");
        PostMethod method = uploadFileParts("fileName", filePartObject);
        assertEquals("<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?><message><message>abc.txt: Bob</message></message>", method.getResponseBodyAsString());
    }

    @Test
    public void uploadWithEncodedNameShouldLeaveFileNameCorrect() throws Exception {
        FilePartObject filePartObject = new FilePartObject("file", false, "=?utf-8?B?0YLQtdGB0YI=?=", "text/plain", "Hello first world!");
        FilePartObject result = (FilePartObject) uploadAndUnmarshalFileParts("single", filePartObject);
        assertEquals("тест", result.getFilename());
    }

    @Test
    public void uploadWithIllegalCharacterEncodingShouldResultInClientError() throws Exception {
        PostMethod method = uploadFileParts("single", new FilePartObject("file", false, "=?utf-unknown?B?unknown?=", "text/plain", "Hello first world!"));
        assertEquals(Response.Status.BAD_REQUEST.getStatusCode(), method.getStatusCode());
    }

    private static void assertFilePartEquals(FilePartObject o1, FilePartObject o2) {
        assertEquals("Field name", o1.getFieldName(), o2.getFieldName());
        assertEquals("Is field", o1.isField(), o2.isField());
        assertEquals("File name", o1.getFilename(), o2.getFilename());
        assertEquals("File contents", o1.getValue(), o2.getValue());
    }

    private PostMethod uploadFileParts(String to, FilePartObject... fileParts) throws Exception {
        ArrayList<Part> parts = new ArrayList<Part>();
        for (FilePartObject o : fileParts) {
            if (o.isField()) {
                StringPart stringPart = new StringPart(o.getFieldName(), o.getValue());
                parts.add(stringPart);
            } else {
                PartSource partSource = new ByteArrayPartSource(o.getFilename(), o.getValue().getBytes("UTF-8"));
                FilePart filePart = new FilePart(o.getFieldName(), partSource, o.getContentType(), "UTF-8");
                parts.add(filePart);
            }
        }

        PostMethod method = new PostMethod(WebResourceFactory.getUriBuilder().path("rest").path("refimpl").path(WebResourceFactory.REST_VERSION).path("multipart").path(to).build().toString() + ".xml");
        method.setRequestHeader(XsrfResourceFilter.TOKEN_HEADER, XsrfResourceFilter.NO_CHECK);
        MultipartRequestEntity requestEntity = new MultipartRequestEntity(parts.toArray(new Part[parts.size()]), new HttpMethodParams());
        method.setRequestEntity(requestEntity);
        client.executeMethod(method);
        return method;
    }

    private Object uploadAndUnmarshalFileParts(String to, FilePartObject... fileParts) throws Exception {
        PostMethod method = uploadFileParts(to, fileParts);
        Unmarshaller unmarshaller = JAXBContext.newInstance(FilePartObject.class, FilePartObjects.class).createUnmarshaller();
        return unmarshaller.unmarshal(method.getResponseBodyAsStream());
    }
}