diff --git a/src/test/java/com/onfido/integration/ReportSchemasTest.java b/src/test/java/com/onfido/integration/ReportSchemasTest.java
index 5278313..25e258c 100644
--- a/src/test/java/com/onfido/integration/ReportSchemasTest.java
+++ b/src/test/java/com/onfido/integration/ReportSchemasTest.java
@@ -27,29 +27,34 @@ public void schemaOfDocumentReportIsValid() throws Exception {
             document,
             new CheckBuilder().reportNames(Arrays.asList(ReportName.DOCUMENT)));
 
-    DocumentReport documentReport =
-        (DocumentReport)
+    Report documentReport =
+        (Report)
             repeatRequestUntilStatusChanges(
                 "findReport", new Object[] {check.getReportIds().get(0)}, COMPLETE, 10, 1000);
 
     Assertions.assertEquals(COMPLETE, documentReport.getStatus());
     Assertions.assertEquals(ReportName.DOCUMENT, documentReport.getName());
-    Assertions.assertEquals(check.getId(), documentReport.getCheckId());
-    Assertions.assertNotNull(documentReport.getBreakdown());
-    Assertions.assertNotNull(documentReport.getProperties());
-    Assertions.assertNotNull(documentReport.getResult());
-    Assertions.assertNotNull(documentReport.getSubResult());
+    Assertions.assertEquals(check.getId(), documentReport.getReportShared().getCheckId());
+    Assertions.assertNotNull(documentReport.getDocumentReport().getBreakdown());
+    Assertions.assertNotNull(documentReport.getDocumentReport().getProperties());
+    Assertions.assertNotNull(documentReport.getReportShared().getResult());
+    Assertions.assertNotNull(documentReport.getReportShared().getSubResult());
 
     Assertions.assertEquals(
         "clear",
         documentReport
+            .getDocumentReport()
             .getBreakdown()
             .getDataComparison()
             .getBreakdown()
             .getIssuingCountry()
             .getResult());
     Assertions.assertEquals(
-        LocalDate.parse("1990-01-01"), documentReport.getProperties().getDateOfBirth());
+        LocalDate.parse("1990-01-01"),
+        documentReport.getDocumentReport().getProperties().getDateOfBirth());
+    Assertions.assertEquals(check.getId(), documentReport.getReportShared().getCheckId());
+    Assertions.assertNotNull(documentReport.getDocumentReport().getBreakdown());
+    Assertions.assertNotNull(documentReport.getDocumentReport().getProperties());
   }
 
   @Test
@@ -63,8 +68,8 @@ public void schemaOfFacialSimilarityPhotoReportIsValid() throws Exception {
             new CheckBuilder()
                 .reportNames(Arrays.asList(ReportName.FACIAL_SIMILARITY_PHOTO_FULLY_AUTO)));
 
-    FacialSimilarityPhotoFullyAutoReport facialSimilarityPhotoFullyAutoReport =
-        (FacialSimilarityPhotoFullyAutoReport)
+    Report facialSimilarityPhotoFullyAutoReport =
+        (Report)
             repeatRequestUntilStatusChanges(
                 "findReport", new Object[] {check.getReportIds().get(0)}, COMPLETE, 10, 1000);
 
@@ -72,9 +77,16 @@ public void schemaOfFacialSimilarityPhotoReportIsValid() throws Exception {
     Assertions.assertEquals(
         ReportName.FACIAL_SIMILARITY_PHOTO_FULLY_AUTO,
         facialSimilarityPhotoFullyAutoReport.getName());
-    Assertions.assertEquals(check.getId(), facialSimilarityPhotoFullyAutoReport.getCheckId());
-    Assertions.assertNotNull(facialSimilarityPhotoFullyAutoReport.getBreakdown());
-    Assertions.assertNotNull(facialSimilarityPhotoFullyAutoReport.getResult());
-    Assertions.assertNotNull(facialSimilarityPhotoFullyAutoReport.getProperties());
+    Assertions.assertEquals(
+        check.getId(), facialSimilarityPhotoFullyAutoReport.getReportShared().getCheckId());
+    Assertions.assertNotNull(
+        facialSimilarityPhotoFullyAutoReport
+            .getFacialSimilarityPhotoFullyAutoReport()
+            .getBreakdown());
+    Assertions.assertNotNull(facialSimilarityPhotoFullyAutoReport.getReportShared().getResult());
+    Assertions.assertNotNull(
+        facialSimilarityPhotoFullyAutoReport
+            .getFacialSimilarityPhotoFullyAutoReport()
+            .getProperties());
   }
 }
diff --git a/src/test/java/com/onfido/integration/TestBase.java b/src/test/java/com/onfido/integration/TestBase.java
index c2eec2a..3db7dcf 100644
--- a/src/test/java/com/onfido/integration/TestBase.java
+++ b/src/test/java/com/onfido/integration/TestBase.java
@@ -13,7 +13,6 @@
 import com.onfido.model.IdPhoto;
 import com.onfido.model.LivePhoto;
 import com.onfido.model.LocationBuilder;
-import com.onfido.model.Report;
 import com.onfido.model.WatchlistMonitor;
 import com.onfido.model.WatchlistMonitorBuilder;
 import com.onfido.model.Webhook;
@@ -175,27 +174,18 @@ public Object repeatRequestUntilStatusChanges(
           InvocationTargetException {
     Method method = getMethod(methodName, params);
     Object instance = method.invoke(onfido, params);
-    boolean isInstanceOfReport = instance instanceof Report;
-
-    if (isInstanceOfReport) {
-      instance = ((Report) instance).getActualInstance();
-    }
 
     int iteration = 0;
+    Method getStatusMethod = instance.getClass().getMethod("getStatus");
 
-    while (!instance.getClass().getMethod("getStatus").invoke(instance).equals(status)) {
+    while (!getStatusMethod.invoke(instance).equals(status)) {
       if (iteration > maxRetries) {
         throw new RuntimeException("Status did not change in time");
       }
 
       iteration += 1;
       Thread.sleep(sleepTime);
-
-      if (isInstanceOfReport) {
-        instance = ((Report) method.invoke(onfido, params)).getActualInstance();
-      } else {
-        instance = method.invoke(onfido, params);
-      }
+      instance = method.invoke(onfido, params);
     }
     return instance;
   }
diff --git a/src/test/java/com/onfido/integration/WorkflowRunTest.java b/src/test/java/com/onfido/integration/WorkflowRunTest.java
index d2338b9..93e510c 100644
--- a/src/test/java/com/onfido/integration/WorkflowRunTest.java
+++ b/src/test/java/com/onfido/integration/WorkflowRunTest.java
@@ -17,6 +17,7 @@ public class WorkflowRunTest extends TestBase {
   private UUID workflowRunId;
 
   static final UUID WORKFLOW_ID = UUID.fromString("e8c921eb-0495-44fe-b655-bcdcaffdafe5");
+  static final UUID WORKFLOW_ID_TIMELINE = UUID.fromString("221f9d24-cf72-4762-ac4a-01bf3ccc09dd");
 
   @BeforeEach
   public void setup() throws Exception {
@@ -69,8 +70,7 @@ public void evidenceWorkflowRunTest() throws Exception {
 
   @Test
   public void generateTimelineFileTest() throws Exception {
-    UUID workflowId = UUID.fromString("221f9d24-cf72-4762-ac4a-01bf3ccc09dd");
-    UUID workflowRunId = createWorkflowRun(workflowId, applicantId).getId();
+    UUID workflowRunId = createWorkflowRun(WORKFLOW_ID_TIMELINE, applicantId).getId();
 
     repeatRequestUntilStatusChanges(
         "findWorkflowRun", new UUID[] {workflowRunId}, WorkflowRun.StatusEnum.APPROVED, 10, 1000);
@@ -82,8 +82,7 @@ public void generateTimelineFileTest() throws Exception {
 
   @Test
   public void findTimelineFileTest() throws Exception {
-    UUID workflowId = UUID.fromString("221f9d24-cf72-4762-ac4a-01bf3ccc09dd");
-    UUID workflowRunId = createWorkflowRun(workflowId, applicantId).getId();
+    UUID workflowRunId = createWorkflowRun(WORKFLOW_ID_TIMELINE, applicantId).getId();
 
     repeatRequestUntilStatusChanges(
         "findWorkflowRun", new UUID[] {workflowRunId}, WorkflowRun.StatusEnum.APPROVED, 10, 1000);