diff --git a/src/Repository/Repository.php b/src/Repository/Repository.php
index 9729655..07ac608 100755
--- a/src/Repository/Repository.php
+++ b/src/Repository/Repository.php
@@ -215,9 +215,6 @@ final protected function mapPropertiesToColumn(array $arguments): array
 
     final protected function createModel(array $data): object
     {
-        /**
-         * @var class-string<AsEntity> $entityName
-         */
         $entityName = $this->getEntityName();
         $primaryKeyValue = $data[$entityName::getPrimaryKeyColumn()];
         $cacheKey = $entityName.$primaryKeyValue;
@@ -228,14 +225,14 @@ final protected function createModel(array $data): object
             $this->cache->set($entity->_getKey(), $entity);
         }
 
-        if ($entity instanceof AsEntity) {
+        if (method_exists($entity, 'hydrate')) {
             $entity->hydrate($data);
         }
 
-        /**
-         * @var object$entity
-         */
-        $entity->setEntityManager($this->manager);
+        if (method_exists($entity, 'setEntityManager')) {
+            $entity->setEntityManager($this->manager);
+        }
+        
         return $entity;
     }
 
diff --git a/tests/Repository/CommentRepository.php b/tests/Repository/CommentRepository.php
deleted file mode 100644
index 647969e..0000000
--- a/tests/Repository/CommentRepository.php
+++ /dev/null
@@ -1,20 +0,0 @@
-<?php
-
-namespace Test\AlphaSoft\AsLinkOrm\Repository;
-
-
-use AlphaSoft\AsLinkOrm\Repository\Repository;
-use Model\Comment;
-
-class CommentRepository extends Repository
-{
-    public function getTableName(): string
-    {
-        return 'comment';
-    }
-
-    public function getEntityName(): string
-    {
-        return Comment::class;
-    }
-}