diff --git a/composer.json b/composer.json
index 66949218..25c8ad02 100644
--- a/composer.json
+++ b/composer.json
@@ -38,7 +38,6 @@
     "psr-4": {
       "RonasIT\\Support\\Tests\\": "tests/",
       "RonasIT\\Support\\Tests\\Support\\": "tests/Support/",
-      "App\\Models\\": "tests/Support/Models/",
       "App\\Nova\\": "tests/Support/Nova/"
     },
     "files": [
diff --git a/src/Generators/AbstractTestsGenerator.php b/src/Generators/AbstractTestsGenerator.php
index 780f889d..46dc0e1c 100644
--- a/src/Generators/AbstractTestsGenerator.php
+++ b/src/Generators/AbstractTestsGenerator.php
@@ -8,7 +8,6 @@
 use Illuminate\Support\Str;
 use RonasIT\Support\Events\SuccessCreateMessage;
 use RonasIT\Support\Exceptions\CircularRelationsFoundedException;
-use RonasIT\Support\Exceptions\ClassNotExistsException;
 
 abstract class AbstractTestsGenerator extends EntityGenerator
 {
@@ -162,11 +161,6 @@ protected function buildEntityObject($model): array
         return $result;
     }
 
-    protected function getModelClass($model): string
-    {
-        return "App\\Models\\{$model}";
-    }
-
     protected function getModelFields($model): array
     {
         $modelClass = $this->getModelClass($model);
@@ -219,7 +213,7 @@ protected function generateFixture($fixtureName, $data): void
     protected function buildRelationsTree($models): array
     {
         foreach ($models as $model) {
-            $relations = $this->getRelatedModels($model);
+            $relations = $this->getRelatedModels($model, $this->getTestClassName());
             $relationsWithFactories = $this->getModelsWithFactories($relations);
 
             if (empty($relationsWithFactories)) {
@@ -242,30 +236,6 @@ protected function buildRelationsTree($models): array
         return array_unique($models);
     }
 
-    protected function getRelatedModels($model)
-    {
-        $content = $this->getModelClassContent($model);
-
-        preg_match_all('/(?<=belongsTo\().*(?=::class)/', $content, $matches);
-
-        return head($matches);
-    }
-
-    protected function getModelClassContent($model): string
-    {
-        $path = base_path("{$this->paths['models']}/{$model}.php");
-
-        if (!$this->classExists('models', $model)) {
-            $this->throwFailureException(
-                ClassNotExistsException::class,
-                "Cannot create {$model} Model cause {$model} Model does not exists.",
-                "Create a {$model} Model by himself or run command 'php artisan make:entity {$model} --only-model'."
-            );
-        }
-
-        return file_get_contents($path);
-    }
-
     protected function canGenerateUserData(): bool
     {
         return $this->classExists('models', 'User')
diff --git a/src/Generators/EntityGenerator.php b/src/Generators/EntityGenerator.php
index 7d62c35a..937c11a4 100644
--- a/src/Generators/EntityGenerator.php
+++ b/src/Generators/EntityGenerator.php
@@ -2,9 +2,15 @@
 
 namespace RonasIT\Support\Generators;
 
+use Illuminate\Database\Eloquent\Relations\BelongsTo;
 use Illuminate\Filesystem\Filesystem;
 use Illuminate\Support\Arr;
+use Illuminate\Support\Facades\DB;
 use Illuminate\Support\Str;
+use RonasIT\Support\Exceptions\ClassNotExistsException;
+use Throwable;
+use ReflectionMethod;
+use ReflectionClass;
 
 /**
  * @property Filesystem $fs
@@ -157,4 +163,52 @@ protected function throwFailureException($exceptionClass, $failureMessage, $reco
     {
         throw new $exceptionClass("{$failureMessage} {$recommendedMessage}");
     }
+
+    protected function getRelatedModels(string $model, string $creatableClass): array
+    {
+        $modelClass = $this->getModelClass($model);
+
+        if (!class_exists($modelClass)) {
+            $this->throwFailureException(
+                exceptionClass: ClassNotExistsException::class,
+                failureMessage: "Cannot create {$creatableClass} cause {$model} Model does not exists.",
+                recommendedMessage: "Create a {$model} Model by himself or run command 'php artisan make:entity {$model} --only-model'.",
+            );
+        }
+
+        $instance = new $modelClass();
+
+        $publicMethods = (new ReflectionClass($modelClass))->getMethods(ReflectionMethod::IS_PUBLIC);
+
+        $methods = array_filter($publicMethods, fn ($method) => $method->class === $modelClass && !$method->getParameters());
+
+        $relatedModels = [];
+
+        DB::beginTransaction();
+
+        foreach ($methods as $method) {
+            try {
+                $result = call_user_func([$instance, $method->getName()]);
+
+                if (!$result instanceof BelongsTo) {
+                    continue;
+                }
+            } catch (Throwable) {
+                continue;
+            }
+
+            $relatedModels[] = class_basename(get_class($result->getRelated()));
+        }
+
+        DB::rollBack();
+
+        return $relatedModels;
+    }
+
+    protected function getModelClass(string $model): string
+    {
+        $modelNamespace = $this->getOrCreateNamespace('models');
+
+        return "{$modelNamespace}\\{$model}";
+    }
 }
diff --git a/src/Generators/FactoryGenerator.php b/src/Generators/FactoryGenerator.php
index f0a9b6a4..5ac02bb7 100644
--- a/src/Generators/FactoryGenerator.php
+++ b/src/Generators/FactoryGenerator.php
@@ -112,7 +112,7 @@ protected function prepareEmptyFactory(): void
     protected function checkExistRelatedModelsFactories(): bool
     {
         $modelFactoryContent = file_get_contents($this->paths['factory']);
-        $relatedModels = $this->getRelatedModels($this->model);
+        $relatedModels = $this->getRelatedModels($this->model, "{$this->model}Factory");
         $modelNamespace = $this->getOrCreateNamespace('models');
 
         foreach ($relatedModels as $relatedModel) {
@@ -238,35 +238,4 @@ protected function getFactoryPattern($model): string
 
         return "/{$modelNamespace}.*{$return}/sU";
     }
-
-    protected function getModelClass($model): string
-    {
-        $modelNamespace = $this->getOrCreateNamespace('models');
-
-        return "{$modelNamespace}\\{$model}";
-    }
-
-    protected function getRelatedModels($model)
-    {
-        $content = $this->getModelClassContent($model);
-
-        preg_match_all('/(?<=belongsTo\().*(?=::class)/', $content, $matches);
-
-        return head($matches);
-    }
-
-    protected function getModelClassContent($model): string
-    {
-        $path = base_path("{$this->paths['models']}/{$model}.php");
-
-        if (!$this->classExists('models', $model)) {
-            $this->throwFailureException(
-                ClassNotExistsException::class,
-                "Cannot create {$model} Model cause {$model} Model does not exists.",
-                "Create a {$model} Model by himself or run command 'php artisan make:entity {$model} --only-model'."
-            );
-        }
-
-        return file_get_contents($path);
-    }
 }
diff --git a/tests/NovaTestGeneratorTest.php b/tests/NovaTestGeneratorTest.php
index 3159b47e..c1ebeccc 100644
--- a/tests/NovaTestGeneratorTest.php
+++ b/tests/NovaTestGeneratorTest.php
@@ -2,12 +2,15 @@
 
 namespace RonasIT\Support\Tests;
 
+use RonasIT\Support\Tests\Support\Models\WelcomeBonus;
 use Illuminate\Support\Facades\Event;
 use RonasIT\Support\Events\SuccessCreateMessage;
 use RonasIT\Support\Exceptions\ClassAlreadyExistsException;
 use RonasIT\Support\Exceptions\ClassNotExistsException;
 use RonasIT\Support\Generators\NovaTestGenerator;
 use RonasIT\Support\Tests\Support\NovaTestGeneratorTest\NovaTestGeneratorMockTrait;
+use Laravel\Nova\NovaServiceProvider;
+use Mockery;
 
 class NovaTestGeneratorTest extends TestCase
 {
@@ -54,7 +57,19 @@ className: ClassAlreadyExistsException::class,
 
     public function testSuccess()
     {
-        $this->mockNovaServiceProviderExists();
+        config([
+            'entity-generator.paths.models' => 'RonasIT/Support/Tests/Support/Models',
+        ]);
+
+        $mock = Mockery::mock('alias:Illuminate\Support\Facades\DB');
+        $mock
+            ->shouldReceive('beginTransaction', 'rollBack')
+            ->once();
+
+        $this->mockNativeGeneratorFunctions(
+            $this->nativeClassExistsMethodCall([NovaServiceProvider::class, true]),
+            $this->nativeClassExistsMethodCall([WelcomeBonus::class, true]),
+        );
 
         $this->mockFilesystem();
         $this->mockNovaRequestClassCall();
diff --git a/tests/Support/GeneratorMockTrait.php b/tests/Support/GeneratorMockTrait.php
index f40e718b..5933ed8d 100644
--- a/tests/Support/GeneratorMockTrait.php
+++ b/tests/Support/GeneratorMockTrait.php
@@ -6,20 +6,16 @@
 
 trait GeneratorMockTrait
 {
-    public function mockClassExistsFunction(string $className, bool $result = true, bool $autoloadArg = true): void
+    public function mockNativeGeneratorFunctions(...$functionCalls): void
     {
-        $this->mockNativeFunction('\RonasIT\Support\Generators', [
-            $this->functionCall(
-                name: 'class_exists',
-                arguments: [$className, $autoloadArg],
-                result: $result,
-            ),
-        ]);
+        $this->mockNativeFunction('\RonasIT\Support\Generators', $functionCalls);
     }
 
     public function mockNovaServiceProviderExists(bool $result = true): void
     {
-        $this->mockClassExistsFunction(NovaServiceProvider::class, $result);
+        $this->mockNativeGeneratorFunctions(
+            $this->nativeClassExistsMethodCall([NovaServiceProvider::class, true], $result),
+        );
     }
 
     public function classExistsMethodCall(array $arguments, bool $result = true): array
@@ -31,6 +27,15 @@ public function classExistsMethodCall(array $arguments, bool $result = true): ar
         ];
     }
 
+    public function nativeClassExistsMethodCall(array $arguments, bool $result = true): array
+    {
+        return [
+            'function' => 'class_exists',
+            'arguments' => $arguments,
+            'result' => $result,
+        ];
+    }
+
     public function mockPhpFileContent(): string
     {
         return '<?php';
diff --git a/tests/Support/Models/WelcomeBonus.php b/tests/Support/Models/WelcomeBonus.php
index bf20ee9b..dcc84688 100644
--- a/tests/Support/Models/WelcomeBonus.php
+++ b/tests/Support/Models/WelcomeBonus.php
@@ -1,6 +1,6 @@
 <?php
 
-namespace App\Models;
+namespace RonasIT\Support\Tests\Support\Models;
 
 use Illuminate\Database\Eloquent\Model;
 use RonasIT\Support\Traits\ModelTrait;