diff --git a/tests/src/Unit/TimeBasedFieldUpdaterTest.php b/tests/src/Unit/TimeBasedFieldUpdaterTest.php
index b97254a..2bee827 100644
--- a/tests/src/Unit/TimeBasedFieldUpdaterTest.php
+++ b/tests/src/Unit/TimeBasedFieldUpdaterTest.php
@@ -17,24 +17,24 @@ class TimeBasedFieldUpdaterTest extends UnitTestCase{
 
     }
     function testFieldIsUpdatedOnEntityReturnedByQuery():void {
-        $this->query->method('execute')->willReturn([1]);
+        $this->query->method('execute')->willReturn(['1']);
         $this->entity = $this->getMockBuilder('\\Drupal\\Core\\Entity\\EntityBase')->disableOriginalConstructor()->getMock();
         $this->storage->expects($this->once())->method('load')->with(1)->willReturn($this->entity);
         $this->updater->update([], '3 months', ['dummy_field' => 4]);
         $this->assertEquals($this->entity->dummy_field, 4);
     }
     function testFieldIsUpdatedOnAllEntitiesReturnedByQuery():void {
-        $this->query->method('execute')->willReturn([1, 2]);
+        $this->query->method('execute')->willReturn(['1', '2']);
         $entity = $this->getEntity();
         $entity2 = $this->getEntity();
-        $this->storage->method('load')->willReturnMap([[1, $entity], [2, $entity2]]);
+        $this->storage->method('load')->willReturnMap([['1', $entity], ['2', $entity2]]);
         $this->updater->update([], '3 months', ['dummy_field' => 4]);
         $this->assertEquals($entity->dummy_field, 4);
         $this->assertEquals($entity2->dummy_field, 4);
     }
 
     function testMultipleFieldsAreUpdated(): void {
-        $this->query->method('execute')->willReturn([1]);
+        $this->query->method('execute')->willReturn(['1']);
         $entity = $this->getEntity();       
         $this->storage->expects($this->once())->method('load')->with(1)->willReturn($entity);
         $this->updater->update([], '3 months', ['dummy_field' => 4, 'dummy_field_2' => 5]);