diff --git a/src/Extensions/ElasticSynonymExtension.php b/src/Extensions/ElasticSynonymExtension.php
index 3899c40..8fe3cd7 100644
--- a/src/Extensions/ElasticSynonymExtension.php
+++ b/src/Extensions/ElasticSynonymExtension.php
@@ -16,6 +16,7 @@
 use Firesphere\ElasticSearch\Models\SynonymSet;
 use Firesphere\ElasticSearch\Services\ElasticCoreService;
 use Firesphere\SearchBackend\Models\SearchSynonym;
+use Psr\Container\NotFoundExceptionInterface;
 use SilverStripe\Core\Injector\Injector;
 use SilverStripe\Forms\FieldList;
 use SilverStripe\ORM\DataExtension;
@@ -27,30 +28,24 @@
  */
 class ElasticSynonymExtension extends DataExtension
 {
-    public function updateCMSFields(FieldList $fields)
-    {
-        $fields->removeByName('SynonymSets');
-        parent::updateCMSFields($fields);
-    }
-
     /**
      * Add or update this synonym in Elastic
      *
-     * @throws ClientResponseException
-     * @throws ServerResponseException
-     * @throws MissingParameterException
+     * @throws NotFoundExceptionInterface
      */
     public function onAfterWrite()
     {
         $service = Injector::inst()->get(ElasticCoreService::class);
+        /** @var SearchSynonym|ElasticSynonymExtension $owner */
+        $owner = $this->owner;
         $syn = $service->getClient()->synonyms();
         /** @var SynonymSet $set */
         $set = SynonymSet::get()->first();
         $syn->putSynonymRule([
             'set_id'  => $set->Key,
-            'rule_id' => $this->owner->getModifiedID(),
+            'rule_id' => $owner->getModifiedID(),
             'body'    => [
-                'synonyms' => $this->owner->getCombinedSynonym()
+                'synonyms' => $owner->getCombinedSynonym()
             ]
         ]);
     }
@@ -58,17 +53,17 @@ public function onAfterWrite()
     /**
      * When deleting a synonym from the CMS, delete it as a rule
      *
-     * @throws ClientResponseException
-     * @throws ServerResponseException
-     * @throws MissingParameterException
+     * @throws NotFoundExceptionInterface
      */
     public function onAfterDelete()
     {
         $service = Injector::inst()->get(ElasticCoreService::class);
         $syn = $service->getClient()->synonyms();
+        /** @var SearchSynonym $owner */
+        $owner = $this->owner;
         /** @var SynonymSet $set */
         $set = SynonymSet::get()->first();
-        $syn->deleteSynonymRule(['set_id' => $set->Key, 'rule_id' => $this->owner->getModifiedId()]);
+        $syn->deleteSynonymRule(['set_id' => $set->Key, 'rule_id' => $owner->getModifiedId()]);
         parent::onAfterDelete();
     }
 }
diff --git a/src/Tasks/ElasticIndexTask.php b/src/Tasks/ElasticIndexTask.php
index 4ca51be..b0ffa81 100644
--- a/src/Tasks/ElasticIndexTask.php
+++ b/src/Tasks/ElasticIndexTask.php
@@ -88,6 +88,11 @@ class ElasticIndexTask extends BuildTask
      */
     protected $index;
 
+    /**
+     * @var int
+     */
+    protected $groups = 0;
+
     /**
      * ElasticIndexTask constructor. Sets up the document factory
      *
@@ -142,7 +147,7 @@ public function run($request)
         $time = gmdate('H:i:s', (time() - $start));
         $this->getLogger()->info(sprintf('Time taken: %s', $time));
 
-        return $groups;
+        $this->groups = $groups;
     }
 
     /**
@@ -217,6 +222,16 @@ public function setDebug(bool $debug, bool $force = false): self
         return $this;
     }
 
+    public function getGroups(): int
+    {
+        return $this->groups;
+    }
+
+    public function setGroups(int $groups): void
+    {
+        $this->groups = $groups;
+    }
+
     /**
      * Index a single class for a given index. {@link static::indexClassForIndex()}
      *
diff --git a/tests/unit/Tasks/ElasticIndexTaskTest.php b/tests/unit/Tasks/ElasticIndexTaskTest.php
index 6653864..790c706 100644
--- a/tests/unit/Tasks/ElasticIndexTaskTest.php
+++ b/tests/unit/Tasks/ElasticIndexTaskTest.php
@@ -33,14 +33,14 @@ public function testRun()
         $page->publishSingle();
         $task = new ElasticIndexTask();
         $request = new HTTPRequest('GET', 'dev/tasks/ElasticIndexTask');
-        $result = $task->run($request);
+        $task->run($request);
 
-        $this->assertGreaterThan(0, $result);
+        $this->assertGreaterThan(0, $task->getGroups());
         $this->assertinstanceOf(ElasticIndex::class, $task->getIndex());
         $request = new HTTPRequest('GET', 'dev/tasks/ElasticIndexTask', ['clear' => true]);
-        $result = $task->run($request);
+        $task->run($request);
 
-        $this->assertGreaterThan(0, $result);
+        $this->assertGreaterThan(0, $task->getGroups());
         $this->assertinstanceOf(ElasticIndex::class, $task->getIndex());
     }
 }