diff --git a/src/Listener/OptimisticLock.php b/src/Listener/OptimisticLock.php
index c0dac10..87e87cd 100644
--- a/src/Listener/OptimisticLock.php
+++ b/src/Listener/OptimisticLock.php
@@ -5,7 +5,7 @@
 namespace Cycle\ORM\Entity\Behavior\Listener;
 
 use Cycle\ORM\Command\ScopeCarrierInterface;
-use Cycle\ORM\Command\Special\WrappedCommand;
+use Cycle\ORM\Command\Special\WrappedStoreCommand;
 use Cycle\ORM\Command\StoreCommandInterface;
 use Cycle\ORM\Entity\Behavior\Attribute\Listen;
 use Cycle\ORM\Entity\Behavior\Event\Mapper\Command\OnCreate;
@@ -76,7 +76,7 @@ public function __invoke(OnDelete|OnUpdate $event): void
         $event->command = $this->lock($event->node, $event->state, $event->command);
     }
 
-    private function lock(Node $node, State $state, ScopeCarrierInterface $command): WrappedCommand
+    private function lock(Node $node, State $state, ScopeCarrierInterface $command): WrappedStoreCommand
     {
         $nodeValue = $node->getData()[$this->field] ?? null;
         if ($nodeValue === null) {
@@ -100,7 +100,7 @@ private function lock(Node $node, State $state, ScopeCarrierInterface $command):
 
         $command->setScope($this->field, $nodeValue);
 
-        return WrappedCommand::wrapCommand($command)
+        return WrappedStoreCommand::wrapCommand($command)
             ->withAfterExecution(static function (ScopeCarrierInterface $command) use ($node): void {
                 if ($command->getAffectedRows() === 0) {
                     throw new RecordIsLockedException($node);
diff --git a/tests/Behavior/Fixtures/OptimisticLock/Author.php b/tests/Behavior/Fixtures/OptimisticLock/Author.php
new file mode 100644
index 0000000..8c61d6d
--- /dev/null
+++ b/tests/Behavior/Fixtures/OptimisticLock/Author.php
@@ -0,0 +1,23 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Cycle\ORM\Entity\Behavior\Tests\Fixtures\OptimisticLock;
+
+use Cycle\Annotated\Annotation\Column;
+use Cycle\Annotated\Annotation\Embeddable;
+
+#[Embeddable]
+class Author
+{
+    public function __construct(
+        #[Column(type: 'string', name: 'author_first_name')]
+        public string $firstName,
+
+        #[Column(type: 'string', name: 'author_last_name')]
+        public string $lastName,
+    )
+    {
+        
+    }
+}
diff --git a/tests/Behavior/Fixtures/OptimisticLock/Comment.php b/tests/Behavior/Fixtures/OptimisticLock/Comment.php
index d5d3205..2d8346c 100644
--- a/tests/Behavior/Fixtures/OptimisticLock/Comment.php
+++ b/tests/Behavior/Fixtures/OptimisticLock/Comment.php
@@ -6,6 +6,7 @@
 
 use Cycle\Annotated\Annotation\Column;
 use Cycle\Annotated\Annotation\Entity;
+use Cycle\Annotated\Annotation\Relation\Embedded;
 use Cycle\ORM\Entity\Behavior\OptimisticLock;
 
 #[Entity]
@@ -15,10 +16,18 @@ class Comment
     #[Column(type: 'primary')]
     public int $id;
 
+    #[Embedded(target: Author::class)]
+    public Author $author;
+
     public ?string $content = null;
     public ?int $versionInt = null;
     public ?string $versionStr = null;
     public ?\DateTimeImmutable $versionDatetime = null;
     public ?string $versionMicrotime = null;
     public ?int $versionCustom = null;
+
+    public function __construct()
+    {
+        $this->author = new Author(firstName: 'First', lastName: 'Last');
+    }
 }
diff --git a/tests/Behavior/Functional/Driver/Common/BaseSchemaTest.php b/tests/Behavior/Functional/Driver/Common/BaseSchemaTest.php
index cb592c6..98e7358 100644
--- a/tests/Behavior/Functional/Driver/Common/BaseSchemaTest.php
+++ b/tests/Behavior/Functional/Driver/Common/BaseSchemaTest.php
@@ -5,6 +5,7 @@
 namespace Cycle\ORM\Entity\Behavior\Tests\Functional\Driver\Common;
 
 use Cycle\Annotated\Entities;
+use Cycle\Annotated\Embeddings;
 use Cycle\Annotated\MergeColumns;
 use Cycle\Annotated\MergeIndexes;
 use Cycle\ORM\Schema;
@@ -31,8 +32,11 @@ public function compileWithTokenizer(Tokenizer $tokenizer): void
     {
         $reader = new AttributeReader();
 
+        $classLocator = $tokenizer->classLocator();
+
         $this->schema = new Schema((new Compiler())->compile($this->registry = new Registry($this->dbal), [
-            new Entities($tokenizer->classLocator(), $reader),
+            new Embeddings($classLocator, $reader),
+            new Entities($classLocator, $reader),
             new ResetTables(),
             new MergeColumns($reader),
             new MergeIndexes($reader),
diff --git a/tests/Behavior/Functional/Driver/Common/OptimisticLock/ListenerTest.php b/tests/Behavior/Functional/Driver/Common/OptimisticLock/ListenerTest.php
index 260732a..240551b 100644
--- a/tests/Behavior/Functional/Driver/Common/OptimisticLock/ListenerTest.php
+++ b/tests/Behavior/Functional/Driver/Common/OptimisticLock/ListenerTest.php
@@ -8,11 +8,13 @@
 use Cycle\ORM\Entity\Behavior\Exception\OptimisticLock\RecordIsLockedException;
 use Cycle\ORM\Entity\Behavior\Listener\OptimisticLock;
 use Cycle\ORM\Entity\Behavior\Tests\Fixtures\OptimisticLock\Comment;
+use Cycle\ORM\Entity\Behavior\Tests\Fixtures\OptimisticLock\Author;
 use Cycle\ORM\Entity\Behavior\Tests\Functional\Driver\Common\BaseListenerTest;
 use Cycle\ORM\Entity\Behavior\Tests\Traits\TableTrait;
 use Cycle\ORM\Heap\Heap;
 use Cycle\ORM\Schema;
 use Cycle\ORM\SchemaInterface;
+use Cycle\ORM\Relation;
 use Cycle\ORM\Select;
 use Cycle\ORM\Transaction;
 
@@ -34,6 +36,8 @@ public function setUp(): void
                 'version_microtime' => 'string',
                 'version_custom' => 'int,nullable',
                 'content' => 'string,nullable',
+                'author_first_name' => 'string',
+                'author_last_name' => 'string',
             ]
         );
 
@@ -81,8 +85,25 @@ public function setUp(): void
                     'versionDatetime' => 'datetime'
                 ],
                 SchemaInterface::SCHEMA => [],
-                SchemaInterface::RELATIONS => [],
+                SchemaInterface::RELATIONS => [
+                    'author' => [
+                        Relation::TYPE => Relation::EMBEDDED,
+                        Relation::TARGET => Author::class,
+                        Relation::LOAD => Relation::LOAD_EAGER,
+                        Relation::SCHEMA => [],
+                    ]
+                ],
             ],
+            Author::class => [
+                SchemaInterface::ENTITY => Author::class,
+                SchemaInterface::DATABASE => 'default',
+                SchemaInterface::TABLE => 'comments',
+                SchemaInterface::PRIMARY_KEY => ['id'],
+                SchemaInterface::COLUMNS => [
+                    'first_name' => 'author_first_name',
+                    'last_name' => 'author_last_name',
+                ],
+            ]
         ]));
     }
 
@@ -95,7 +116,7 @@ public function testAddVersionOnCreate()
 
         $this->orm = $this->orm->with(heap: new Heap());
         $select = new Select($this->orm, Comment::class);
-
+        
         $comment = $select->fetchOne();
 
         $this->assertSame(1, $comment->versionInt);