diff --git a/lib/API/JSON/v1/DynamicSegments.php b/lib/API/JSON/v1/DynamicSegments.php
index 28415dc831..e70e489d9a 100644
--- a/lib/API/JSON/v1/DynamicSegments.php
+++ b/lib/API/JSON/v1/DynamicSegments.php
@@ -120,49 +120,45 @@ class DynamicSegments extends APIEndpoint {
   }
 
   public function trash($data = []) {
-    if (isset($data['id'])) {
-      $id = (int)$data['id'];
-    } else {
+    if (!isset($data['id'])) {
       return $this->errorResponse([
         Error::BAD_REQUEST => WPFunctions::get()->__('Missing mandatory argument `id`.', 'mailpoet'),
       ]);
     }
 
-    try {
-      $segment = $this->dynamicSegmentsLoader->load($id);
-      $segment->trash();
-      return $this->successResponse(
-        $segment->asArray(),
-        ['count' => 1]
-      );
-    } catch (\InvalidArgumentException $e) {
+    $segment = $this->getSegment($data);
+    if ($segment === null) {
       return $this->errorResponse([
         Error::NOT_FOUND => WPFunctions::get()->__('This segment does not exist.', 'mailpoet'),
       ]);
     }
+
+    $this->segmentsRepository->bulkTrash([$segment->getId()], SegmentEntity::TYPE_DYNAMIC);
+    return $this->successResponse(
+      $this->segmentsResponseBuilder->build($segment),
+      ['count' => 1]
+    );
   }
 
   public function restore($data = []) {
-    if (isset($data['id'])) {
-      $id = (int)$data['id'];
-    } else {
+    if (!isset($data['id'])) {
       return $this->errorResponse([
         Error::BAD_REQUEST => WPFunctions::get()->__('Missing mandatory argument `id`.', 'mailpoet'),
       ]);
     }
 
-    try {
-      $segment = $this->dynamicSegmentsLoader->load($id);
-      $segment->restore();
-      return $this->successResponse(
-        $segment->asArray(),
-        ['count' => 1]
-      );
-    } catch (\InvalidArgumentException $e) {
+    $segment = $this->getSegment($data);
+    if ($segment === null) {
       return $this->errorResponse([
         Error::NOT_FOUND => WPFunctions::get()->__('This segment does not exist.', 'mailpoet'),
       ]);
     }
+
+    $this->segmentsRepository->bulkRestore([$segment->getId()], SegmentEntity::TYPE_DYNAMIC);
+    return $this->successResponse(
+      $this->segmentsResponseBuilder->build($segment),
+      ['count' => 1]
+    );
   }
 
   public function delete($data = []) {
@@ -210,4 +206,10 @@ class DynamicSegments extends APIEndpoint {
       ]);
     }
   }
+
+  private function getSegment(array $data): ?SegmentEntity {
+    return isset($data['id'])
+      ? $this->segmentsRepository->findOneById((int)$data['id'])
+      : null;
+  }
 }
diff --git a/tests/integration/API/JSON/v1/DynamicSegmentsTest.php b/tests/integration/API/JSON/v1/DynamicSegmentsTest.php
index 8f6d3a9d71..7427715b31 100644
--- a/tests/integration/API/JSON/v1/DynamicSegmentsTest.php
+++ b/tests/integration/API/JSON/v1/DynamicSegmentsTest.php
@@ -107,55 +107,33 @@ class DynamicSegmentsTest extends \MailPoetTest {
   }
 
   public function testItCanTrashASegment() {
-    DynamicSegment::deleteMany();
-    $dynamicSegment = DynamicSegment::createOrUpdate([
-      'name' => 'Trash test',
-      'description' => 'description',
-    ]);
-    $loader = Stub::makeEmpty('\MailPoet\DynamicSegments\Persistence\Loading\SingleSegmentLoader', [
-      'load' => function () use($dynamicSegment) {
-        return $dynamicSegment;
-      },
-    ]);
+    $dynamicSegment = $this->createDynamicSegmentEntity('Trash test', 'description');
 
-    $endpoint = new DynamicSegments($this->bulkAction, $this->listingHandler, $this->listingRepository, $this->responseBuilder, $this->segmentsRepository, $this->saveController, $loader);
-    $response = $endpoint->trash(['id' => $dynamicSegment->id]);
+    $endpoint = new DynamicSegments($this->bulkAction, $this->listingHandler, $this->listingRepository, $this->responseBuilder, $this->segmentsRepository, $this->saveController);
+    $response = $endpoint->trash(['id' => $dynamicSegment->getId()]);
 
     expect($response->status)->equals(self::SUCCESS_RESPONSE_CODE);
-    expect($response->data)->equals($dynamicSegment->asArray());
+    expect($response->data['name'])->equals($dynamicSegment->getName());
     expect($response->meta['count'])->equals(1);
 
-    $dynamicSegment = DynamicSegment::findOne($dynamicSegment->id);
-    assert($dynamicSegment instanceof DynamicSegment);
-    expect($dynamicSegment->deletedAt)->notNull();
-
-    $dynamicSegment->delete();
+    $this->entityManager->refresh($dynamicSegment);
+    assert($dynamicSegment instanceof SegmentEntity);
+    expect($dynamicSegment->getDeletedAt())->notNull();
   }
 
   public function testItCanRestoreASegment() {
-    DynamicSegment::deleteMany();
-    $dynamicSegment = DynamicSegment::createOrUpdate([
-      'name' => 'Restore test',
-      'description' => 'description',
-    ]);
-    $loader = Stub::makeEmpty('\MailPoet\DynamicSegments\Persistence\Loading\SingleSegmentLoader', [
-      'load' => function () use($dynamicSegment) {
-        return $dynamicSegment;
-      },
-    ]);
+    $dynamicSegment = $this->createDynamicSegmentEntity('Trash test', 'description');
 
-    $endpoint = new DynamicSegments($this->bulkAction, $this->listingHandler, $this->listingRepository, $this->responseBuilder, $this->segmentsRepository, $this->saveController, $loader);
-    $response = $endpoint->restore(['id' => $dynamicSegment->id]);
+    $endpoint = new DynamicSegments($this->bulkAction, $this->listingHandler, $this->listingRepository, $this->responseBuilder, $this->segmentsRepository, $this->saveController);
+    $response = $endpoint->restore(['id' => $dynamicSegment->getId()]);
 
     expect($response->status)->equals(self::SUCCESS_RESPONSE_CODE);
-    expect($response->data)->equals($dynamicSegment->asArray());
+    expect($response->data['name'])->equals($dynamicSegment->getName());
     expect($response->meta['count'])->equals(1);
 
-    $dynamicSegment = DynamicSegment::findOne($dynamicSegment->id);
-    assert($dynamicSegment instanceof DynamicSegment);
-    expect($dynamicSegment->deletedAt)->equals(null);
-
-    $dynamicSegment->delete();
+    $this->entityManager->refresh($dynamicSegment);
+    assert($dynamicSegment instanceof SegmentEntity);
+    expect($dynamicSegment->getDeletedAt())->null();
   }
 
   public function testItCanDeleteASegment() {