diff --git a/lib/Segments/SegmentSaveController.php b/lib/Segments/SegmentSaveController.php
index b9ea587eca..30fb200ea0 100644
--- a/lib/Segments/SegmentSaveController.php
+++ b/lib/Segments/SegmentSaveController.php
@@ -4,7 +4,6 @@ namespace MailPoet\Segments;
 
 use InvalidArgumentException;
 use MailPoet\Entities\SegmentEntity;
-use MailPoet\NotFoundException;
 
 class SegmentSaveController {
   /** @var SegmentsRepository */
@@ -19,8 +18,11 @@ class SegmentSaveController {
   public function save(array $data = []): SegmentEntity {
     $this->checkSegmenUniqueName($data['name'] ?? '');
 
-    $segment = $this->createOrUpdateSegment($data);
-    return $segment;
+    $id = isset($data['id']) ? (int)$data['id'] : null;
+    $name = $data['name'] ?? '';
+    $description = $data['description'] ?? '';
+
+    return $this->segmentsRepository->createOrUpdate($name, $description, $id);
   }
 
   private function checkSegmenUniqueName(string $name): void {
@@ -29,34 +31,4 @@ class SegmentSaveController {
       throw new InvalidArgumentException("Segment with name: '{$name}' already exists.");
     }
   }
-
-  private function createOrUpdateSegment(array $data): SegmentEntity {
-    $name = $data['name'] ?? '';
-    $description = $data['description'] ?? '';
-
-    if (isset($data['id'])) {
-      $segment = $this->getSegment((int)$data['id']);
-      $segment->setName($name);
-      $segment->setDescription($description);
-    } else {
-      $segment = new SegmentEntity(
-        $name,
-        SegmentEntity::TYPE_DEFAULT,
-        $description
-      );
-      $this->segmentsRepository->persist($segment);
-    }
-
-    $this->segmentsRepository->flush();
-    return $segment;
-  }
-
-  private function getSegment(int $segmentId): SegmentEntity {
-    $segment = $this->segmentsRepository->findOneById($segmentId);
-    if (!$segment) {
-      throw new NotFoundException();
-    }
-
-    return $segment;
-  }
 }
diff --git a/lib/Segments/SegmentsRepository.php b/lib/Segments/SegmentsRepository.php
index 8807695d4b..a8fbf76dd3 100644
--- a/lib/Segments/SegmentsRepository.php
+++ b/lib/Segments/SegmentsRepository.php
@@ -4,6 +4,7 @@ namespace MailPoet\Segments;
 
 use MailPoet\Doctrine\Repository;
 use MailPoet\Entities\SegmentEntity;
+use MailPoet\NotFoundException;
 
 /**
  * @extends Repository<SegmentEntity>
@@ -31,4 +32,24 @@ class SegmentsRepository extends Repository {
     }
     return $countMap;
   }
+
+  public function createOrUpdate(
+    string $name,
+    string $description = '',
+    ?int $id = null
+  ): SegmentEntity {
+    if ($id) {
+      $segment = $this->findOneById($id);
+      if (!$segment instanceof SegmentEntity) {
+        throw new NotFoundException("Segment with ID [{$id}] was not found.");
+      }
+      $segment->setName($name);
+      $segment->setDescription($description);
+    } else {
+      $segment = new SegmentEntity($name, SegmentEntity::TYPE_DEFAULT, $description);
+      $this->persist($segment);
+    }
+    $this->flush();
+    return $segment;
+  }
 }