diff --git a/Test/Unit/Model/Api/GuestAdyenStateDataTest.php b/Test/Unit/Model/Api/GuestAdyenStateDataTest.php
index d2a9110a9..29809c43d 100644
--- a/Test/Unit/Model/Api/GuestAdyenStateDataTest.php
+++ b/Test/Unit/Model/Api/GuestAdyenStateDataTest.php
@@ -14,6 +14,7 @@
 use Adyen\Payment\Helper\StateData;
 use Adyen\Payment\Model\Api\GuestAdyenStateData;
 use Adyen\Payment\Test\Unit\AbstractAdyenTestCase;
+use Magento\Framework\App\State;
 use Magento\Framework\Exception\InputException;
 use Magento\Framework\TestFramework\Unit\Helper\ObjectManager;
 use Magento\Quote\Model\QuoteIdMask;
@@ -51,10 +52,16 @@ public function testSaveSuccessful()
         $stateData = '{"stateData":"dummyData"}';
         $cartId = 'ABC123456789';
 
+        $stateDataMock = $this->createMock(\Adyen\Payment\Model\StateData::class);
+        $stateDataMock->method('getEntityId')->willReturn(1);
+
         $this->quoteIdMaskFactoryMock->method('create')->willReturn($this->quoteIdMaskMock);
-        $this->stateDataHelperMock->expects($this->once())->method('saveStateData');
+        $this->stateDataHelperMock->expects($this->once())->method('saveStateData')->willReturn($stateDataMock);
+        //should return int
+
 
         $this->guestAdyenStateDataModel->save($stateData, $cartId);
+//        $this->assertIsInt($result);
     }
 
     public function testRemoveSuccessful()
@@ -64,7 +71,6 @@ public function testRemoveSuccessful()
 
         $this->quoteIdMaskFactoryMock->method('create')->willReturn($this->quoteIdMaskMock);
         $this->stateDataHelperMock->expects($this->once())->method('removeStateData');
-
         $this->guestAdyenStateDataModel->remove($stateDataId, $cartId);
     }