diff --git a/src/InternalReplicator.php b/src/InternalReplicator.php
index 65d8c93..b897ffa 100644
--- a/src/InternalReplicator.php
+++ b/src/InternalReplicator.php
@@ -93,6 +93,7 @@ public function replicate(WorkspacePointerInterface $source, WorkspacePointerInt
     $source_workspace = $source->getWorkspace();
     $target_workspace = $target->getWorkspace();
     // Set active workspace to source.
+    $original_workspace = $this->workspaceManager->getActiveWorkspace();
     try {
       $this->workspaceManager->setActiveWorkspace($source_workspace);
     }
@@ -103,46 +104,19 @@ public function replicate(WorkspacePointerInterface $source, WorkspacePointerInt
     // Fetch the site time.
     $start_time = new \DateTime();
 
-    // If no task sent, create an empty task for its defaults.
-    if ($task === NULL) {
-      $task = new ReplicationTask();
-    }
-
-    // Get changes on the source workspace.
-    $source_changes = $this->changesFactory->get($source_workspace)
-        ->filter($task->getFilter())
-        ->parameters($task->getParameters())
-        ->getNormal();
-    $data = [];
-    foreach ($source_changes as $source_change) {
-      $data[$source_change['id']] = [];
-      foreach ($source_change['changes'] as $change) {
-        $data[$source_change['id']][] = $change['rev'];
-      }
-    }
-
-    // Get revisions the target workspace is missing.
-    $revs_diff = $this->revisionDiffFactory->get($target_workspace)->setRevisionIds($data)->getMissing();
-    $entities = [];
-    foreach ($revs_diff as $uuid => $revs) {
-      foreach ($revs['missing'] as $rev) {
-        $missing_found++;
-        $item = $this->revIndex->useWorkspace($source_workspace->id())->get("$uuid:$rev");
-        $entity_type_id = $item['entity_type_id'];
-        $revision_id = $item['revision_id'];
+    $source_data = $this->listSourceChanges($source_workspace, $task);
+    $entities = $this->diffWorkspaces($source_workspace, $target_workspace, $source_data);
 
-        $storage = $this->entityTypeManager->getStorage($entity_type_id);
-        $entity = $storage->loadRevision($revision_id);
-        if ($entity instanceof ContentEntityInterface) {
-          $docs_read++;
-          $entities[] = $this->serializer->normalize($entity, 'json', ['new_revision_id' => TRUE]);
-        }
-      }
+    $normal_entities = [];
+    foreach ($entities as $entity) {
+      $docs_read++;
+      // Convert the entity object to json for BulkDocs
+      $normal_entities[] = $this->serializer->normalize($entity, 'json', ['new_revision_id' => TRUE]);
     }
 
     $data = [
       'new_edits' => FALSE,
-      'docs' => $entities,
+      'docs' => $normal_entities,
     ];
     // Save all entities in bulk.
     $bulk_docs = $this->serializer->denormalize($data, 'Drupal\replication\BulkDocs\BulkDocs', 'json', ['workspace' => $target_workspace]);
@@ -177,7 +151,66 @@ public function replicate(WorkspacePointerInterface $source, WorkspacePointerInt
     $replication_log->setSourceLastSeq($source_workspace->getUpdateSeq());
     $replication_log->setHistory($history);
     $replication_log->save();
+
+    // return to original active workspace
+    try {
+      $this->workspaceManager->setActiveWorkspace($original_workspace);
+    }
+    catch (\Exception $e) {
+      watchdog_exception('Workspace', $e);
+      drupal_set_message($e->getMessage(), 'error');
+    }
+
     return $replication_log;
   }
 
+  /**
+   * List the changes on the source according to the task
+   */
+  public function listSourceChanges(WorkspaceInterface $source_workspace, ReplicationTaskInterface $task = NULL) {
+    // If no task sent, create an empty task for its defaults.
+    if ($task === NULL) {
+      $task = new ReplicationTask();
+    }
+
+    // Get changes on the source workspace.
+    $source_changes = $this->changesFactory->get($source_workspace)
+        ->filter($task->getFilter())
+        ->parameters($task->getParameters())
+        ->getNormal();
+    $data = [];
+    foreach ($source_changes as $source_change) {
+      $data[$source_change['id']] = [];
+      foreach ($source_change['changes'] as $change) {
+        $data[$source_change['id']][] = $change['rev'];
+      }
+    }
+
+    return $data;
+  }
+
+  /**
+   * List the entity revisions that need to be replicated
+   */
+  public function diffWorkspaces(WorkspaceInterface $source_workspace, WorkspaceInterface $target_workspace, array $source_data) {
+    // Get revisions the target workspace is missing.
+    $revs_diff = $this->revisionDiffFactory->get($target_workspace)->setRevisionIds($source_data)->getMissing();
+    $entities = [];
+    foreach ($revs_diff as $uuid => $revs) {
+      foreach ($revs['missing'] as $rev) {
+        $missing_found++;
+        $item = $this->revIndex->useWorkspace($source_workspace->id())->get("$uuid:$rev");
+        $entity_type_id = $item['entity_type_id'];
+        $revision_id = $item['revision_id'];
+
+        $storage = $this->entityTypeManager->getStorage($entity_type_id);
+        $entity = $storage->loadRevision($revision_id);
+        if ($entity instanceof ContentEntityInterface) {
+          $entities[] = $entity;
+        }
+      }
+    }
+    return $entities;
+  }
+
 }