From c8a6fed3c6e91a8d44cd21761bfd8e91ed237c6f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Samuel=20=C5=A0tancl?= <samuel.stancl@gmail.com>
Date: Wed, 9 Oct 2019 18:52:38 +0200
Subject: [PATCH] $tenant->run()

---
 src/Tenant.php            | 22 ++++++++++++++
 src/TenantManager.php     |  4 +++
 tests/TenantClassTest.php | 63 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 89 insertions(+)

diff --git a/src/Tenant.php b/src/Tenant.php
index 836e27b8..eb44466f 100644
--- a/src/Tenant.php
+++ b/src/Tenant.php
@@ -5,6 +5,7 @@ declare(strict_types=1);
 namespace Stancl\Tenancy;
 
 use ArrayAccess;
+use Closure;
 use Illuminate\Foundation\Application;
 use Illuminate\Support\Str;
 use Illuminate\Support\Traits\ForwardsCalls;
@@ -362,6 +363,27 @@ class Tenant implements ArrayAccess
         return $this;
     }
 
+    /**
+     * Run a closure inside the tenant's environment.
+     *
+     * @param Closure $closure
+     * @return mixed
+     */
+    public function run(Closure $closure)
+    {
+        $originalTenant = $this->manager->getTenant();
+
+        $this->manager->initializeTenancy($this);
+        $result = $closure($this);
+        $this->manager->endTenancy($this);
+
+        if ($originalTenant) {
+            $this->manager->initializeTenancy($originalTenant);
+        }
+
+        return $result;
+    }
+
     public function __get($key)
     {
         return $this->get($key);
diff --git a/src/TenantManager.php b/src/TenantManager.php
index 2fb5f412..9e2f2c6d 100644
--- a/src/TenantManager.php
+++ b/src/TenantManager.php
@@ -206,6 +206,10 @@ class TenantManager
      */
     public function initializeTenancy(Tenant $tenant): self
     {
+        if ($this->initialized) {
+            $this->endTenancy();
+        }
+
         $this->setTenant($tenant);
         $this->bootstrapTenancy($tenant);
         $this->initialized = true;
diff --git a/tests/TenantClassTest.php b/tests/TenantClassTest.php
index 4a3775f2..a3ee56ca 100644
--- a/tests/TenantClassTest.php
+++ b/tests/TenantClassTest.php
@@ -105,4 +105,67 @@ class TenantClassTest extends TestCase
 
         $this->assertSame(['foo' => 'bar'], $data);
     }
+
+    /** @test */
+    public function run_method_works()
+    {
+        $this->assertSame(null, tenancy()->getTenant());
+
+        $users_table_empty = function () {
+            return count(\DB::table('users')->get()) === 0;
+        };
+
+        $tenant = Tenant::new()->save();
+        \Artisan::call('tenants:migrate', [
+            '--tenants' => [$tenant->id],
+        ]);
+        tenancy()->initialize($tenant);
+        $this->assertTrue($users_table_empty());
+        tenancy()->end();
+
+        $foo = $tenant->run(function () {
+            \DB::table('users')->insert([
+                'name' => 'foo',
+                'email' => 'foo@bar.xy',
+                'password' => bcrypt('secret'),
+            ]);
+
+            return 'foo';
+        });
+
+        // test return value
+        $this->assertSame('foo', $foo);
+
+        // test that tenancy was ended
+        $this->assertSame(false, tenancy()->initialized);
+        $this->assertSame(null, tenancy()->getTenant());
+
+        // test closure
+        tenancy()->initialize($tenant);
+        $this->assertFalse($users_table_empty());
+
+        // test returning to original tenant
+        $tenant2 = Tenant::new()->save();
+        \Artisan::call('tenants:migrate', [
+            '--tenants' => [$tenant2->id],
+        ]);
+
+        tenancy()->initialize($tenant2);
+        $this->assertSame($tenant2, tenancy()->getTenant());
+        $this->assertTrue($users_table_empty());
+
+        $tenant->run(function () {
+            \DB::table('users')->insert([
+                'name' => 'bar',
+                'email' => 'bar@bar.xy',
+                'password' => bcrypt('secret'),
+            ]);
+        });
+
+        $this->assertSame($tenant2, tenancy()->getTenant());
+
+        $this->assertSame(2, $tenant->run(function () {
+            return \DB::table('users')->count();
+        }));
+    }
 }