diff --git a/docgen/pom.xml b/docgen/pom.xml
index 17728230..3bde7f77 100644
--- a/docgen/pom.xml
+++ b/docgen/pom.xml
@@ -1,16 +1,15 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
     <modelVersion>4.0.0</modelVersion>
 
     <parent>
-        <artifactId>error-prone-support</artifactId>
         <groupId>tech.picnic.error-prone-support</groupId>
+        <artifactId>error-prone-support</artifactId>
         <version>0.3.1-SNAPSHOT</version>
     </parent>
 
     <artifactId>docgen</artifactId>
 
-
     <name>Picnic :: Error Prone Support :: Docgen</name>
     <description>Docgen.</description>
 
@@ -28,6 +27,10 @@
             <groupId>${groupId.error-prone}</groupId>
             <artifactId>error_prone_core</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.fasterxml.jackson.core</groupId>
+            <artifactId>jackson-databind</artifactId>
+        </dependency>
         <dependency>
             <groupId>com.google.auto.service</groupId>
             <artifactId>auto-service-annotations</artifactId>
@@ -46,10 +49,5 @@
             <groupId>com.google.guava</groupId>
             <artifactId>guava</artifactId>
         </dependency>
-        <dependency>
-            <groupId>com.fasterxml.jackson.core</groupId>
-            <artifactId>jackson-databind</artifactId>
-        </dependency>
     </dependencies>
-
-</project>
\ No newline at end of file
+</project>
diff --git a/docgen/src/main/java/tech/picnic/errorprone/plugin/BugPatternExtractor.java b/docgen/src/main/java/tech/picnic/errorprone/plugin/BugPatternExtractor.java
index b95345a2..4d3028b4 100644
--- a/docgen/src/main/java/tech/picnic/errorprone/plugin/BugPatternExtractor.java
+++ b/docgen/src/main/java/tech/picnic/errorprone/plugin/BugPatternExtractor.java
@@ -1,18 +1,15 @@
 package tech.picnic.errorprone.plugin;
 
-//public final class BugPatternExtractor implements DocExtractor<BugPatternData> {
-//  @Override
-//  public BugPatternTestData extractData(ClassTree tree, VisitorState state) {
-//    String name = tree.getSimpleName().toString().replace("Test", "");
-//    ScanBugCheckerTestData scanner = new ScanBugCheckerTestData(state);
-//
-//    tree.getMembers().stream()
-//        .filter(MethodTree.class::isInstance)
-//        .map(MethodTree.class::cast)
-//        .filter(m -> BUG_PATTERN_TEST.matches(m, state))
-//        .forEach(m -> scanner.scan(m, null));
-//
-//    return BugPatternTestData.create(
-//        name, scanner.getIdentification(), scanner.getInput(), scanner.getOutput());
-//  }
-//}
+import com.google.errorprone.BugPattern;
+import com.google.errorprone.VisitorState;
+import com.sun.source.tree.ClassTree;
+import com.sun.source.util.TaskEvent;
+import tech.picnic.errorprone.plugin.objects.BugPatternData;
+
+public final class BugPatternExtractor implements DocExtractor<BugPatternData> {
+  @Override
+  public BugPatternData extractData(ClassTree tree, TaskEvent taskEvent, VisitorState state) {
+    BugPattern annotation = taskEvent.getTypeElement().getAnnotation(BugPattern.class);
+    return BugPatternData.create(annotation, taskEvent.getTypeElement().getSimpleName().toString());
+  }
+}
diff --git a/docgen/src/main/java/tech/picnic/errorprone/plugin/BugPatternTestsExtractor.java b/docgen/src/main/java/tech/picnic/errorprone/plugin/BugPatternTestsExtractor.java
index f11a7a2a..4abdde12 100644
--- a/docgen/src/main/java/tech/picnic/errorprone/plugin/BugPatternTestsExtractor.java
+++ b/docgen/src/main/java/tech/picnic/errorprone/plugin/BugPatternTestsExtractor.java
@@ -13,6 +13,7 @@ import com.sun.source.tree.ClassTree;
 import com.sun.source.tree.ExpressionTree;
 import com.sun.source.tree.MethodInvocationTree;
 import com.sun.source.tree.MethodTree;
+import com.sun.source.util.TaskEvent;
 import com.sun.source.util.TreeScanner;
 import java.util.List;
 import javax.annotation.Nullable;
@@ -37,7 +38,7 @@ public final class BugPatternTestsExtractor implements DocExtractor<BugPatternTe
           .named("addOutputLines");
 
   @Override
-  public BugPatternTestData extractData(ClassTree tree, VisitorState state) {
+  public BugPatternTestData extractData(ClassTree tree, TaskEvent taskEvent, VisitorState state) {
     String name = tree.getSimpleName().toString().replace("Test", "");
     ScanBugCheckerTestData scanner = new ScanBugCheckerTestData(state);
 
diff --git a/docgen/src/main/java/tech/picnic/errorprone/plugin/DocExtractor.java b/docgen/src/main/java/tech/picnic/errorprone/plugin/DocExtractor.java
index dc1da436..06421b8e 100644
--- a/docgen/src/main/java/tech/picnic/errorprone/plugin/DocExtractor.java
+++ b/docgen/src/main/java/tech/picnic/errorprone/plugin/DocExtractor.java
@@ -2,7 +2,8 @@ package tech.picnic.errorprone.plugin;
 
 import com.google.errorprone.VisitorState;
 import com.sun.source.tree.ClassTree;
+import com.sun.source.util.TaskEvent;
 
 public interface DocExtractor<T> {
-  T extractData(ClassTree tree, VisitorState state);
+  T extractData(ClassTree tree, TaskEvent taskEvent, VisitorState state);
 }
diff --git a/docgen/src/main/java/tech/picnic/errorprone/plugin/DocgenTaskListener.java b/docgen/src/main/java/tech/picnic/errorprone/plugin/DocgenTaskListener.java
index b0aa56aa..c0d13ab8 100644
--- a/docgen/src/main/java/tech/picnic/errorprone/plugin/DocgenTaskListener.java
+++ b/docgen/src/main/java/tech/picnic/errorprone/plugin/DocgenTaskListener.java
@@ -52,17 +52,18 @@ final class DocgenTaskListener implements TaskListener {
     }
 
     if (isBugPatternTest(tree)) {
-      BugPatternTestData testData = new BugPatternTestsExtractor().extractData(tree, state);
+      BugPatternTestData testData =
+          new BugPatternTestsExtractor().extractData(tree, taskEvent, state);
       System.out.println(testData);
       writeToFile(testData, "bug-pattern-test-data.jsonl");
-    } else {
+    } else if(isBugPattern(tree)) {
 
-      BugPattern annotation = taskEvent.getTypeElement().getAnnotation(BugPattern.class);
-      BugPatternData bugPatternData =
-          BugPatternData.create(annotation, taskEvent.getTypeElement().getSimpleName().toString());
+      BugPatternData data = new BugPatternExtractor().extractData(tree, taskEvent, state);
 
       System.out.println("Analysing: " + taskEvent.getTypeElement().getSimpleName());
-      writeToFile(bugPatternData, "bug-pattern-data.jsonl");
+      writeToFile(data, "bug-pattern-data.jsonl");
+    } else {
+      System.out.println("~~~~~~~~~!!!!~~~~~~~NOW ANALYSING: " + tree.getSimpleName().toString());
     }
   }
 
diff --git a/error-prone-contrib/pom.xml b/error-prone-contrib/pom.xml
index e8681ecd..dbbf20a2 100644
--- a/error-prone-contrib/pom.xml
+++ b/error-prone-contrib/pom.xml
@@ -238,4 +238,32 @@
             </plugins>
         </pluginManagement>
     </build>
+
+    <profiles>
+        <profile>
+            <id>docgen</id>
+            <build>
+                <plugins>
+                    <plugin>
+                        <groupId>org.codehaus.mojo</groupId>
+                        <artifactId>build-helper-maven-plugin</artifactId>
+                        <executions>
+                            <execution>
+                                <id>add-test-source</id>
+                                <goals>
+                                    <goal>add-test-source</goal>
+                                </goals>
+                                <phase>generate-test-sources</phase>
+                                <configuration>
+                                    <sources>
+                                        <source>src/test/resources</source>
+                                    </sources>
+                                </configuration>
+                            </execution>
+                        </executions>
+                    </plugin>
+                </plugins>
+            </build>
+        </profile>
+    </profiles>
 </project>