Repository: CleitonDeLima/django-easy-tenants (language: python)

Issue:
Set/Unset tenant in shell context



Discussion / hints:

Fix patch:
diff --git a/easy_tenants/conf.py b/easy_tenants/conf.py
index b8333cc..6a58409 100644
--- a/easy_tenants/conf.py
+++ b/easy_tenants/conf.py
@@ -8,6 +8,7 @@
 
 class EasyTenantsAppConf(AppConf):
     TENANT_FIELD = "tenant"
+    TENANT_MODEL = ""
 
     class Meta:
         prefix = "EASY_TENANTS"
diff --git a/easy_tenants/management/__init__.py b/easy_tenants/management/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/easy_tenants/management/commands/__init__.py b/easy_tenants/management/commands/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/easy_tenants/management/commands/shell_tenant.py b/easy_tenants/management/commands/shell_tenant.py
new file mode 100644
index 0000000..5aea0c3
--- /dev/null
+++ b/easy_tenants/management/commands/shell_tenant.py
@@ -0,0 +1,56 @@
+from django.apps import apps
+from django.conf import settings
+from django.core.management import call_command
+from django.core.management.base import BaseCommand
+
+from easy_tenants import tenant_context, tenant_context_disabled
+
+try:
+    from django_extensions.management.commands import shell_plus as shell_cmd
+except ImportError:
+    from django.core.management.commands import shell as shell_cmd
+
+
+class Command(BaseCommand):
+    help = (
+        "Run a Python REPL scoped to a specific tenant. "
+        "Specify the --tenant_id or uses --disable to ignore tenant context."
+    )
+    tests_mode = False
+
+    def add_arguments(self, parser):
+        self.cmd = shell_cmd.Command()
+        self.cmd.add_arguments(parser)
+
+        parser.add_argument("--tenant_id")
+        parser.add_argument(
+            "--disable",
+            action="store_true",
+            help="All tenant are disabled for shell session.",
+        )
+
+    def handle(self, *args, **options):
+        tenant_id = options.pop("tenant_id", None)
+        disable = options.pop("disable", False)
+
+        if disable:
+            with tenant_context_disabled():
+                self.stdout.write(
+                    self.style.SUCCESS(
+                        "All tenant are disabled for this shell session "
+                        "– please be careful!"
+                    )
+                )
+                return self.call_command(*args, **options)
+
+        TenantModel = apps.get_model(settings.EASY_TENANTS_TENANT_MODEL)
+        tenant = TenantModel.objects.get(id=tenant_id)
+
+        with tenant_context(tenant):
+            return self.call_command(*args, **options)
+
+    def call_command(self, *args, **options):
+        if self.tests_mode:
+            return "shell"
+
+        return call_command(self.cmd, *args, **options)
diff --git a/example/settings.py b/example/settings.py
index 2dffe1b..b64d936 100644
--- a/example/settings.py
+++ b/example/settings.py
@@ -131,3 +131,4 @@
 DEFAULT_AUTO_FIELD = "django.db.models.AutoField"
 
 EASY_TENANTS_TENANT_FIELD = "tenant"
+EASY_TENANTS_TENANT_MODEL = "app_test.Customer"


Test patch:
diff --git a/tests/settings.py b/tests/settings.py
index c673899..b00e01b 100644
--- a/tests/settings.py
+++ b/tests/settings.py
@@ -46,3 +46,4 @@
 STATIC_URL = "/static/"
 
 EASY_TENANTS_TENANT_FIELD = "store"
+EASY_TENANTS_TENANT_MODEL = "tests.StoreTenant"
diff --git a/tests/test_commands.py b/tests/test_commands.py
new file mode 100644
index 0000000..0f4818d
--- /dev/null
+++ b/tests/test_commands.py
@@ -0,0 +1,29 @@
+from io import StringIO
+
+import pytest
+from django.core.management import call_command
+
+from easy_tenants.management.commands import shell_tenant
+from tests.models import StoreTenant
+
+
+@pytest.mark.django_db
+class TestShellTenant:
+    def test_set_tenant_context(self):
+        tenant = StoreTenant.objects.create()
+
+        cmd = shell_tenant.Command()
+        cmd.tests_mode = True
+
+        assert call_command(cmd, tenant_id=str(tenant.id)) == "shell"
+
+    def test_override(self):
+        stdout = StringIO()
+        cmd = shell_tenant.Command()
+        cmd.tests_mode = True
+
+        call_command(cmd, "--disable", stdout=stdout)
+        assert (
+            "All tenant are disabled for this shell session – "
+            "please be careful!"
+        ) in stdout.getvalue()