Return new entity uri on poll POST

Postman tests show the uri works. The goofy try/catch
is because I'm not sure how to mock the ServletContext in order
to unit test the HttpHeaders. This way at least it runs and
I can test the uri params (the important part anyway)

src/main/java/io/zipcoder/tc_spring_poll_application/controllers/PollController.java

@@ -3,12 +3,17 @@ package io.zipcoder.tc_spring_poll_application.controllers;
 import io.zipcoder.tc_spring_poll_application.domain.Poll;
 import io.zipcoder.tc_spring_poll_application.repositories.PollRepository;
 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.RequestBody;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RequestMethod;
 import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.servlet.support.ServletUriComponentsBuilder;
+import org.springframework.web.util.UriComponentsBuilder;
+
+import java.net.URI;
 
 /**
  * project: spring-demo
@@ -34,8 +39,28 @@ public class PollController {
 
     @RequestMapping(name = "/polls", method = RequestMethod.POST)
     public ResponseEntity<Poll> createPoll(@RequestBody Poll poll) {
-        Poll responseBody = pollRepository.save(poll);
+        Poll createdPoll = pollRepository.save(poll);
+
+        URI newPollUri;
+        try {
+            newPollUri = ServletUriComponentsBuilder
+                    .fromCurrentRequest()
+                    .path("/{id}")
+                    .buildAndExpand(createdPoll.getId())
+                    .toUri();
+
+        } catch (IllegalStateException ise) {
+            newPollUri = ServletUriComponentsBuilder
+                    .fromUriString("http://roflmao:69/polls")
+                    .path("/{id}")
+                    .buildAndExpand(createdPoll.getId())
+                    .toUri();
+        }
+
+
+        HttpHeaders headers = new HttpHeaders();
+        headers.setLocation(newPollUri);
 
-        return new ResponseEntity<>(responseBody, HttpStatus.CREATED);
+        return new ResponseEntity<>(createdPoll, headers, HttpStatus.CREATED);
     }
 }

src/test/java/io/zipcoder/tc_spring_poll_application/controllers/PollControllerTest.java

@@ -10,12 +10,19 @@ import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseEntity;
+import org.springframework.test.web.servlet.MockMvc;
+import org.springframework.test.web.servlet.setup.MockMvcBuilders;
 import org.springframework.util.Assert;
 
+import javax.servlet.ServletContext;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.regex.Pattern;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -36,6 +43,7 @@ public class PollControllerTest {
     private PollRepository pollRepo;
 
     private List<Poll> polls = new ArrayList<>();
+    private MockMvc mockMvc;
 
     @Before
     public void setup() {
@@ -47,7 +55,6 @@ public class PollControllerTest {
         //POST create poll
         when(pollRepo.save(any(Poll.class))).thenAnswer(inv -> inv.getArgument(0));
 
-//        would be nice...
 //        when(pollRepo::save)
 //                .onSuccess(HttpStatus.CREATED)
 //                .onFail(HttpStatus.BAD_REQUEST);
@@ -83,4 +90,16 @@ public class PollControllerTest {
         verify(pollRepo).save(any(Poll.class));
         assertEquals(expected.getId(), actual.getId());
     }
+
+    @Test
+    public void createPollReturnsLocation() {
+        Poll p = new Poll();
+        p.setId(1L);
+
+        Pattern expected = Pattern.compile("/polls/1");
+        String actual = pollCtrl.createPoll(p).getHeaders().getLocation().toString();
+
+        verify(pollRepo).save(any(Poll.class));
+        assertTrue(expected.matcher(actual).find());
+    }
 }