diff --git a/webgoat-lessons/xxe/src/main/java/org/owasp/webgoat/plugin/Comments.java b/webgoat-lessons/xxe/src/main/java/org/owasp/webgoat/plugin/Comments.java
index 22ba7cf72..8bcf3f618 100644
--- a/webgoat-lessons/xxe/src/main/java/org/owasp/webgoat/plugin/Comments.java
+++ b/webgoat-lessons/xxe/src/main/java/org/owasp/webgoat/plugin/Comments.java
@@ -20,6 +20,10 @@ import java.io.IOException;
 import java.io.StringReader;
 import java.util.Collection;
 import java.util.Map;
+import java.util.Optional;
+
+import static java.util.Optional.empty;
+import static java.util.Optional.of;
 
 /**
  * @author nbaars
@@ -67,12 +71,12 @@ public class Comments {
         return (Comment) unmarshaller.unmarshal(xsr);
     }
 
-    protected Comment parseJson(String comment) {
+    protected Optional<Comment> parseJson(String comment) {
         ObjectMapper mapper = new ObjectMapper();
         try {
-            return mapper.readValue(comment, Comment.class);
+            return of(mapper.readValue(comment, Comment.class));
         } catch (IOException e) {
-            return new Comment();
+            return empty();
         }
     }
 
diff --git a/webgoat-lessons/xxe/src/main/java/org/owasp/webgoat/plugin/ContentTypeAssignment.java b/webgoat-lessons/xxe/src/main/java/org/owasp/webgoat/plugin/ContentTypeAssignment.java
index 96d470c87..6f45e8bf8 100644
--- a/webgoat-lessons/xxe/src/main/java/org/owasp/webgoat/plugin/ContentTypeAssignment.java
+++ b/webgoat-lessons/xxe/src/main/java/org/owasp/webgoat/plugin/ContentTypeAssignment.java
@@ -61,27 +61,26 @@ public class ContentTypeAssignment extends AssignmentEndpoint {
     @ResponseBody
     public AttackResult createNewUser(@RequestBody String commentStr, @RequestHeader("Content-Type") String contentType) throws Exception {
         AttackResult attackResult = failed().build();
-        Comment comment = null;
+
         if (APPLICATION_JSON_VALUE.equals(contentType)) {
-            comment = comments.parseJson(commentStr);
-            comments.addComment(comment, true);
+            comments.parseJson(commentStr).ifPresent(c -> comments.addComment(c, true));
             attackResult = failed().feedback("xxe.content.type.feedback.json").build();
         }
 
         if (MediaType.APPLICATION_XML_VALUE.equals(contentType)) {
             String error = "";
             try {
-                comment = comments.parseXml(commentStr);
+                Comment comment = comments.parseXml(commentStr);
                 comments.addComment(comment, false);
+                if (checkSolution(comment)) {
+                    attackResult = success().build();
+                }
             } catch (Exception e) {
                 error = org.apache.commons.lang.exception.ExceptionUtils.getFullStackTrace(e);
+                attackResult = failed().feedback("xxe.content.type.feedback.xml").output(error).build();
             }
-            attackResult = failed().feedback("xxe.content.type.feedback.xml").output(error).build();
         }
 
-        if (checkSolution(comment)) {
-            attackResult = success().build();
-        }
         return trackProgress(attackResult);
     }
 
@@ -89,7 +88,7 @@ public class ContentTypeAssignment extends AssignmentEndpoint {
         String[] directoriesToCheck = OS.isFamilyUnix() ? DEFAULT_LINUX_DIRECTORIES : DEFAULT_WINDOWS_DIRECTORIES;
         boolean success = true;
         for (String directory : directoriesToCheck) {
-            success &= comment.getText().contains(directory);
+            success &= org.apache.commons.lang3.StringUtils.contains(comment.getText(), directory);
         }
         return success;
     }
diff --git a/webgoat-lessons/xxe/src/test/java/org/owasp/webgoat/plugin/ContentTypeAssignmentTest.java b/webgoat-lessons/xxe/src/test/java/org/owasp/webgoat/plugin/ContentTypeAssignmentTest.java
new file mode 100644
index 000000000..ef3e6b4c1
--- /dev/null
+++ b/webgoat-lessons/xxe/src/test/java/org/owasp/webgoat/plugin/ContentTypeAssignmentTest.java
@@ -0,0 +1,75 @@
+package org.owasp.webgoat.plugin;
+
+import org.hamcrest.CoreMatchers;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.owasp.webgoat.plugins.LessonTest;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.http.MediaType;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+import org.springframework.test.web.servlet.request.MockMvcRequestBuilders;
+import org.springframework.test.web.servlet.setup.MockMvcBuilders;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.when;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
+
+/**
+ * @author nbaars
+ * @since 11/2/17.
+ */
+@RunWith(SpringJUnit4ClassRunner.class)
+public class ContentTypeAssignmentTest extends LessonTest {
+
+    @Autowired
+    private Comments comments;
+
+    @Before
+    public void setup() throws Exception {
+        XXE xxe = new XXE();
+        when(webSession.getCurrentLesson()).thenReturn(xxe);
+        this.mockMvc = MockMvcBuilders.webAppContextSetup(this.wac).build();
+        when(webSession.getUserName()).thenReturn("unit-test");
+    }
+
+    @Test
+    public void sendingXmlButContentTypeIsJson() throws Exception {
+        mockMvc.perform(MockMvcRequestBuilders.post("/xxe/content-type")
+                .contentType(MediaType.APPLICATION_JSON)
+                .content("<?xml version=\"1.0\" standalone=\"yes\" ?><!DOCTYPE user [<!ENTITY root SYSTEM \"file:///\"> ]><comment><text>&root;</text></comment>"))
+                .andExpect(status().isOk())
+                .andExpect(jsonPath("$.feedback", CoreMatchers.is(messages.getMessage("xxe.content.type.feedback.json"))));
+    }
+
+    @Test
+    public void workingAttack() throws Exception {
+        mockMvc.perform(MockMvcRequestBuilders.post("/xxe/content-type")
+                .contentType(MediaType.APPLICATION_XML)
+                .content("<?xml version=\"1.0\" standalone=\"yes\" ?><!DOCTYPE user [<!ENTITY root SYSTEM \"file:///\"> ]><comment><text>&root;</text></comment>"))
+                .andExpect(status().isOk())
+                .andExpect(jsonPath("$.feedback", CoreMatchers.is(messages.getMessage("assignment.solved"))));
+    }
+
+    @Test
+    public void postingJsonShouldAddComment() throws Exception {
+        mockMvc.perform(MockMvcRequestBuilders.post("/xxe/content-type")
+                .contentType(MediaType.APPLICATION_JSON)
+                .content("{  \"text\" : \"Hello World\"}"))
+                .andExpect(status().isOk())
+                .andExpect(jsonPath("$.feedback", CoreMatchers.is(messages.getMessage("xxe.content.type.feedback.json"))));
+        assertThat(comments.getComments().stream().filter(c -> c.getText().equals("Hello World")).count()).isEqualTo(1);
+    }
+
+    @Test
+    public void postingInvalidJsonShouldAddComment() throws Exception {
+        mockMvc.perform(MockMvcRequestBuilders.post("/xxe/content-type")
+                .contentType(MediaType.APPLICATION_JSON)
+                .content("{  'text' : 'Wrong'"))
+                .andExpect(status().isOk())
+                .andExpect(jsonPath("$.feedback", CoreMatchers.is(messages.getMessage("xxe.content.type.feedback.json"))));
+        assertThat(comments.getComments().stream().filter(c -> c.getText().equals("Wrong")).count()).isEqualTo(0);
+    }
+
+}
\ No newline at end of file
diff --git a/webgoat-lessons/xxe/src/test/java/org/owasp/webgoat/plugin/SimpleXXETest.java b/webgoat-lessons/xxe/src/test/java/org/owasp/webgoat/plugin/SimpleXXETest.java
new file mode 100644
index 000000000..ae826e995
--- /dev/null
+++ b/webgoat-lessons/xxe/src/test/java/org/owasp/webgoat/plugin/SimpleXXETest.java
@@ -0,0 +1,66 @@
+package org.owasp.webgoat.plugin;
+
+import org.hamcrest.CoreMatchers;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.owasp.webgoat.plugins.LessonTest;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+import org.springframework.test.web.servlet.request.MockMvcRequestBuilders;
+import org.springframework.test.web.servlet.setup.MockMvcBuilders;
+
+import static org.mockito.Mockito.when;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
+
+/**
+ * @author nbaars
+ * @since 11/2/17.
+ */
+@RunWith(SpringJUnit4ClassRunner.class)
+public class SimpleXXETest extends LessonTest {
+
+    @Before
+    public void setup() throws Exception {
+        XXE xxe = new XXE();
+        when(webSession.getCurrentLesson()).thenReturn(xxe);
+        this.mockMvc = MockMvcBuilders.webAppContextSetup(this.wac).build();
+        when(webSession.getUserName()).thenReturn("unit-test");
+    }
+
+
+    @Test
+    public void workingAttack() throws Exception {
+        //Call with XXE injection
+        mockMvc.perform(MockMvcRequestBuilders.post("/xxe/simple")
+                .content("<?xml version=\"1.0\" standalone=\"yes\" ?><!DOCTYPE user [<!ENTITY root SYSTEM \"file:///\"> ]><comment><text>&root;</text></comment>"))
+                .andExpect(status().isOk())
+                .andExpect(jsonPath("$.feedback", CoreMatchers.is(messages.getMessage("assignment.solved"))));
+    }
+
+    @Test
+    public void postingJsonCommentShouldNotSolveAssignment() throws Exception {
+        mockMvc.perform(MockMvcRequestBuilders.post("/xxe/simple")
+                .content("<comment><text>test</ext></comment>"))
+                .andExpect(status().isOk())
+                .andExpect(jsonPath("$.feedback", CoreMatchers.is(messages.getMessage("assignment.not.solved"))));
+    }
+
+    @Test
+    public void postingXmlCommentWithoutXXEShouldNotSolveAssignment() throws Exception {
+        mockMvc.perform(MockMvcRequestBuilders.post("/xxe/simple")
+                .content("<?xml version=\"1.0\" standalone=\"yes\" ?><comment><text>&root;</text></comment>"))
+                .andExpect(status().isOk())
+                .andExpect(jsonPath("$.feedback", CoreMatchers.is(messages.getMessage("assignment.not.solved"))));
+    }
+
+    @Test
+    public void postingPlainTextShouldShwoException() throws Exception {
+        mockMvc.perform(MockMvcRequestBuilders.post("/xxe/simple")
+                .content("test"))
+                .andExpect(status().isOk())
+                .andExpect(jsonPath("$.output", CoreMatchers.startsWith("javax.xml.bind.UnmarshalException\\n - with linked exception")))
+                .andExpect(jsonPath("$.feedback", CoreMatchers.is(messages.getMessage("assignment.not.solved"))));
+    }
+
+}
\ No newline at end of file