diff --git a/src/FakeConnection.php b/src/FakeConnection.php
index cd5365b..1f2bc0c 100644
--- a/src/FakeConnection.php
+++ b/src/FakeConnection.php
@@ -5,6 +5,7 @@
 use Closure;
 use Illuminate\Database\Connection;
 use Illuminate\Database\ConnectionInterface;
+use Illuminate\Database\Schema\MySqlBuilder;
 use Illuminate\Support\Arr;
 
 class FakeConnection extends Connection implements ConnectionInterface
@@ -27,7 +28,7 @@ public function transaction(Closure $callback, $attempts = 1)
 
     public function getSchemaBuilder()
     {
-        return new FakeSchemaBuilder($this);
+        return new MySqlBuilder($this);
     }
 
     public function getSchemaGrammar()
diff --git a/src/FakeDB.php b/src/FakeDB.php
index d04f65d..47ed8ce 100644
--- a/src/FakeDB.php
+++ b/src/FakeDB.php
@@ -62,7 +62,7 @@ public static function createTable($args)
 
     public static function table($table)
     {
-        return new class ($table){
+        return new class ($table) {
             private $table;
 
             public function __construct($table)
@@ -74,17 +74,22 @@ public function addRow($row)
             {
                 FakeDB::addRow($this->table, $row);
             }
-        };
-    }
 
-    public static function getChangedModel(string $action, $index, $model)
-    {
-        return FakeDB::$changedModels[$model][$action][$index] ?? null;
-    }
+            public function allRows()
+            {
+                $rows = [];
+                foreach((FakeDB::$fakeRows[$this->table] ?? []) as $i => $row) {
+                    $rows[$i] = $row[$this->table];
+                }
 
-    public static function setChangedModel(string $action, $model)
-    {
-        FakeDB::$changedModels[get_class($model)][$action][] = $model;
+                return $rows;
+            }
+
+            public function count()
+            {
+                return count(FakeDB::$fakeRows[$this->table] ?? []);
+            }
+        };
     }
 
     public static function truncate($query = null)
diff --git a/src/FakeEloquentBuilder.php b/src/FakeEloquentBuilder.php
deleted file mode 100644
index 1068b6e..0000000
--- a/src/FakeEloquentBuilder.php
+++ /dev/null
@@ -1,44 +0,0 @@
-<?php
-
-namespace Imanghafoori\EloquentMockery;
-
-use Illuminate\Database\Eloquent\Builder;
-use Illuminate\Database\Eloquent\Model;
-
-class FakeEloquentBuilder extends Builder
-{
-    public function __construct(FakeQueryBuilder $query, Model $modelObj)
-    {
-        $this->query = $query;
-        $this->model = $modelObj;
-    }
-
-    public function delete()
-    {
-        try {
-            return $count = parent::delete();
-        } finally {
-            if (isset($count) && is_int($count) && $count > 0) {
-                FakeDB::setChangedModel('deleted', $this->model);
-            }
-        }
-    }
-
-    public function forceDelete()
-    {
-        try {
-            return $count = parent::forceDelete();
-        } finally {
-            if ($count !== 0) {
-                FakeDB::setChangedModel('deleted', $this->model);
-            }
-        }
-    }
-
-    public function update(array $values)
-    {
-        $this->model->getAttributes() && FakeDB::setChangedModel('updated', $this->model);
-
-        return parent::update($values);
-    }
-}
diff --git a/src/FakeGrammar.php b/src/FakeGrammar.php
index c69b7e3..a328ec3 100644
--- a/src/FakeGrammar.php
+++ b/src/FakeGrammar.php
@@ -88,13 +88,11 @@ public function compileRandom($seed)
 
     public function compileTruncate(Builder $query)
     {
-        return
-            $this->stringy([
-                'type' => 'truncate',
-                'builder' => $query,
-                'sql' => parent::compileTruncate($query),
-            ])
-        ;
+        return $this->stringy([
+            'type' => 'truncate',
+            'builder' => $query,
+            'sql' => parent::compileTruncate($query),
+        ]);
     }
 
     private function stringy($query)
diff --git a/src/FakeSchemaBuilder.php b/src/FakeSchemaBuilder.php
deleted file mode 100644
index 8d85591..0000000
--- a/src/FakeSchemaBuilder.php
+++ /dev/null
@@ -1,10 +0,0 @@
-<?php
-
-namespace Imanghafoori\EloquentMockery;
-
-use Illuminate\Database\Schema\MySqlBuilder;
-
-class FakeSchemaBuilder extends MySqlBuilder
-{
-    //
-}
\ No newline at end of file
diff --git a/src/MockableModel.php b/src/MockableModel.php
deleted file mode 100644
index ce0f5b9..0000000
--- a/src/MockableModel.php
+++ /dev/null
@@ -1,89 +0,0 @@
-<?php
-
-namespace Imanghafoori\EloquentMockery;
-
-trait MockableModel
-{
-    public static function shouldReceive($method)
-    {
-        return new class (self::class, $method){
-
-            private $theClass;
-
-            private $method;
-
-            public function __construct($class, $method)
-            {
-                $this->theClass = $class;
-                $this->method = $method;
-            }
-
-            public function andReturn($value)
-            {
-                $this->theClass::$forceMocks[$this->method][] = $value;
-            }
-        };
-    }
-
-    public static function fakeSoftDelete()
-    {
-        $callback = function ($model) {
-            FakeDB::setChangedModel('softDeleted', $model);
-        };
-        if (method_exists(static::class, 'softDeleted')) {
-            static::softDeleted($callback);
-        } else {
-            static::deleted($callback);
-        }
-    }
-
-    public static function getUpdatedModel($index = 0)
-    {
-        return FakeDB::getChangedModel('updated', $index, static::class);
-    }
-
-    public static function getCreatedModel($index = 0)
-    {
-        return FakeDB::getChangedModel('created', $index, static::class);
-    }
-
-    public static function getSavedModel($index = 0)
-    {
-        return FakeDB::getChangedModel('saved', $index, static::class);
-    }
-
-    public static function getSoftDeletedModel($index = 0)
-    {
-        return FakeDB::getChangedModel('softDeleted', $index, static::class);
-    }
-
-    public static function getDeletedModel($index = 0)
-    {
-        return FakeDB::getChangedModel('deleted', $index, static::class);
-    }
-
-    public static function addFakeRow(array $attributes)
-    {
-        FakeDB::addRow((new static())->getTable(), $attributes);
-    }
-
-    private function isFakeMode()
-    {
-        return isset(FakeDB::$fakeRows[$this->getTable()]);
-    }
-
-    protected function finishSave(array $options)
-    {
-        if (! $this->isFakeMode()) {
-            return parent::finishSave($options);
-        }
-
-        if ($this->wasRecentlyCreated) {
-            FakeDB::setChangedModel('created', $this);
-        }
-
-        FakeDB::setChangedModel('saved', $this);
-
-        return parent::finishSave($options);
-    }
-}
diff --git a/tests/Relations/PolymorphTest.php b/tests/Relations/PolymorphTest.php
new file mode 100644
index 0000000..1b52359
--- /dev/null
+++ b/tests/Relations/PolymorphTest.php
@@ -0,0 +1,217 @@
+<?php
+
+namespace Imanghafoori\EloquentMockery\Tests\Relations;
+
+use Illuminate\Database\Eloquent\Collection;
+use Illuminate\Database\Eloquent\Model;
+use Illuminate\Database\Eloquent\Relations\MorphMany;
+use Illuminate\Database\Eloquent\Relations\MorphOne;
+use Illuminate\Database\Eloquent\Relations\MorphTo;
+use Illuminate\Database\Eloquent\Relations\MorphToMany;
+use Illuminate\Database\Eloquent\Relations\Relation;
+use Imanghafoori\EloquentMockery\FakeDB;
+use PHPUnit\Framework\TestCase;
+
+class PolymorphTest extends TestCase
+{
+    public function tearDown(): void
+    {
+        FakeDB::dontMockQueryBuilder();
+    }
+
+    public function setUp(): void
+    {
+        FakeDB::mockQueryBuilder();
+    }
+
+    /**
+     * @test
+     */
+    public function polymorph_one_to_one_test()
+    {
+        FakeDB::addRow('videos', ['id' => 3, 'name' => 'video_1']);
+        PolymorphVideo::first()->comment()->create(['body' => 'nice video!']);
+
+        $video = PolymorphComment::query()->first()->commentable;
+
+        $this->assertInstanceOf(PolymorphVideo::class, $video);
+        $this->assertEquals(3, $video->id);
+        $this->assertEquals('video_1', $video->name);
+
+        $comment = $video->comment;
+        $this->assertInstanceOf(PolymorphComment::class, $comment);
+        $this->assertEquals(1, $comment->id);
+        $this->assertEquals('nice video!', $comment->body);
+        $this->assertEquals(3, $comment->commentable_id);
+        $this->assertEquals(PolymorphVideo::class, $comment->commentable_type);
+
+        PolymorphVideo::first()->comment()->delete();
+        $this->assertEquals(PolymorphVideo::first()->comment()->count(), 0);
+
+        Relation::morphMap(['video' => PolymorphVideo::class]);
+
+        PolymorphVideo::first()->comment()->create(['body' => 'good video!']);
+
+        $comment = PolymorphVideo::first()->comment()->first();
+
+        $this->assertEquals('good video!', $comment->body);
+        $this->assertEquals(3, $comment->commentable_id);
+        $this->assertEquals('video', $comment->commentable_type);
+
+        $this->assertEquals(PolymorphVideo::first()->comment()->count(), 1);
+
+        Relation::$morphMap = [];
+    }
+
+    /**
+     * @test
+     */
+    public function polymorph_one_to_many()
+    {
+        FakeDB::addRow('videos', ['id' => 3, 'name' => 'video_1']);
+        PolymorphVideo::first()->comments()->create(['body' => 'nice video1']);
+        PolymorphVideo::first()->comments()->create(['body' => 'nice video2']);
+        PolymorphVideo::first()->comments()->create(['body' => 'nice video3']);
+
+        $video = PolymorphComment::query()->first()->commentable;
+
+        $this->assertInstanceOf(PolymorphVideo::class, $video);
+        $this->assertEquals(3, $video->id);
+        $this->assertEquals('video_1', $video->name);
+
+        $comments = $video->comments;
+        $this->assertInstanceOf(Collection::class, $comments);
+        $this->assertCount(3, $comments);
+
+        $comments = $video->comments()->get();
+        $this->assertInstanceOf(Collection::class, $comments);
+        $this->assertCount(3, $comments);
+
+        $comment = $comments[0];
+        $this->assertEquals(1, $comment->id);
+        $this->assertEquals('nice video1', $comment->body);
+        $this->assertEquals(3, $comment->commentable_id);
+        $this->assertEquals(PolymorphVideo::class, $comment->commentable_type);
+
+        $comment = $comments[2];
+        $this->assertEquals(3, $comment->id);
+        $this->assertEquals('nice video3', $comment->body);
+        $this->assertEquals(3, $comment->commentable_id);
+        $this->assertEquals(PolymorphVideo::class, $comment->commentable_type);
+
+        $this->assertEquals(PolymorphVideo::first()->comments()->count(), 3);
+
+        PolymorphVideo::first()->comments()->delete();
+
+        $this->assertEquals(PolymorphVideo::first()->comments()->count(), 0);
+    }
+
+    /**
+     * @test
+     */
+    public function polymorph_many_to_many()
+    {
+        $video = PolymorphVideo::create(['name' => 'vid 1']);
+        $audio = PolymorphAudio::create(['name' => 'aud 1']);
+        $tag = PolymorphTag::create(['name' => 'aud 1']);
+
+        $tag->videos()->attach($video);
+
+        $this->assertCount(1, $tag->videos);
+        $this->assertEquals('vid 1', $tag->videos[0]->name);
+
+        $rows = FakeDB::table('taggables')->allRows();
+
+        $this->assertEquals([
+            'polymorph_tag_id' => 1,
+            'taggable_id' => 1,
+            'taggable_type' => PolymorphVideo::class,
+            'id' => 1,
+        ], $rows[0]);
+        $count = FakeDB::table('taggables')->count();
+        $this->assertEquals(1, $count);
+
+        $tag->audios()->attach($audio);
+
+        $this->assertCount(1, $tag->audios);
+        $this->assertEquals('aud 1', $tag->audios[0]->name);
+
+        $rows = FakeDB::table('taggables')->allRows();
+
+        $count = FakeDB::table('taggables')->count();
+        $this->assertEquals(2, $count);
+
+        $this->assertEquals([
+            'polymorph_tag_id' => 1,
+            'taggable_id' => 1,
+            'taggable_type' => PolymorphVideo::class,
+            'id' => 1,
+        ], $rows[0]);
+
+        $this->assertEquals([
+            'polymorph_tag_id' => 1,
+            'taggable_id' => 1,
+            'taggable_type' => PolymorphAudio::class,
+            'id' => 2,
+        ], $rows[1]);
+    }
+}
+
+class PolymorphComment extends Model
+{
+    public $fillable = ['body', 'commentable_id', 'commentable_type'];
+
+    protected $table = 'comments';
+
+    public function commentable(): MorphTo
+    {
+        return $this->morphTo('commentable', 'commentable_type', 'commentable_id');
+    }
+}
+
+class PolymorphVideo extends Model
+{
+    public $fillable = ['name'];
+
+    protected $table = 'videos';
+
+    public function comment(): MorphOne
+    {
+        return $this->morphOne(PolymorphComment::class, 'commentable');
+    }
+
+    public function comments(): MorphMany
+    {
+        return $this->morphMany(PolymorphComment::class, 'commentable');
+    }
+}
+
+class PolymorphAudio extends Model
+{
+    public $fillable = ['name'];
+
+    protected $table = 'audios';
+
+    public function comments(): MorphMany
+    {
+        return $this->morphMany(PolymorphComment::class, 'commentable');
+    }
+}
+
+class PolymorphTag extends Model
+{
+    public $fillable = ['name'];
+
+    protected $table = 'tags';
+
+    public function audios(): MorphToMany
+    {
+        return $this->morphedByMany(PolymorphAudio::class, 'taggable');
+    }
+
+    public function videos(): MorphToMany
+    {
+        return $this->morphedByMany(PolymorphVideo::class, 'taggable');
+    }
+}
+