diff --git a/database/migrations/2021_05_30_100933_make_documents_assigned_user_nullable.php b/database/migrations/2021_05_30_100933_make_documents_assigned_user_nullable.php
index 6b4f3b3ea5f8..3135b0d5fc74 100644
--- a/database/migrations/2021_05_30_100933_make_documents_assigned_user_nullable.php
+++ b/database/migrations/2021_05_30_100933_make_documents_assigned_user_nullable.php
@@ -1,5 +1,6 @@
 <?php
 
+use App\Libraries\MultiDB;
 use App\Models\Document;
 use Illuminate\Database\Migrations\Migration;
 use Illuminate\Database\Schema\Blueprint;
@@ -21,8 +22,9 @@ class MakeDocumentsAssignedUserNullable extends Migration
         Document::where('assigned_user_id', 0)->update(['assigned_user_id' => null]);
 
         if(config('ninja.db.multi_db_enabled')){
-            Document::on('db-ninja-01')->where('assigned_user_id', 0)->update(['assigned_user_id' => null]);
-            Document::on('db-ninja-02')->where('assigned_user_id', 0)->update(['assigned_user_id' => null]);
+            foreach (MultiDB::$dbs as $db) {
+                Document::on($db)->where('assigned_user_id', 0)->update(['assigned_user_id' => null]);
+            }
         }
         else{
             Document::where('assigned_user_id', 0)->update(['assigned_user_id' => null]);
diff --git a/tests/Feature/Import/ImportCompanyTest.php b/tests/Feature/Import/ImportCompanyTest.php
index d4918b3a1150..c0360ab9e8aa 100644
--- a/tests/Feature/Import/ImportCompanyTest.php
+++ b/tests/Feature/Import/ImportCompanyTest.php
@@ -48,6 +48,7 @@ use App\Models\User;
 use App\Models\Vendor;
 use App\Models\VendorContact;
 use App\Utils\Traits\MakesHash;
+use Illuminate\Foundation\Testing\DatabaseTransactions;
 use Illuminate\Routing\Middleware\ThrottleRequests;
 use Illuminate\Support\Facades\Cache;
 use Illuminate\Support\Str;
@@ -63,7 +64,8 @@ use Tests\TestCase;
 class ImportCompanyTest extends TestCase
 {
     use MakesHash;
-
+    use DatabaseTransactions;
+    
     public $account;
     public $company;
     public $backup_json_object;