Do not touch maintenance mode if already set

When we want to enable/disable maintenance mode manually or externally,
we wouldn't want pyramid_heroku to reset it.

CHANGES.rst

@@ -2,6 +2,14 @@
 Changes
 =======
 
+0.8.0
+-----
+
+ * Do not touch Heroku maintenance mode during migration if it's already enabled.
+   This helps when we want to enable/disable the maintenance mode manually or externally.
+   [sayanarijit]
+
+
 0.7.0
 -----
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pyramid_heroku"
-version = "0.7.0"
+version = "0.8.0"
 description = "A bunch of helpers for successfully running Pyramid on Heroku."
 readme = "README.rst"
 include = [ "CHANGES.rst", "COPYING.txt" ]

pyramid_heroku/__init__.py

@@ -3,7 +3,6 @@
 from ast import literal_eval
 from expandvars import expandvars
 
-import sys
 import typing as t
 
 

pyramid_heroku/migrate.py

@@ -110,6 +110,11 @@ def alembic(self):
         """
         self.shell(f"alembic -c etc/alembic.ini -x ini={self.ini_file} upgrade head")
 
+    def get_maintenance(self) -> bool:
+        res = self.session.get(f"{self.api_endpoint}/apps/{self.app_name}")
+        self.parse_response(res)
+        return res.json()["maintenance"]
+
     def set_maintenance(self, state: bool) -> None:
         res = self.session.patch(
             f"{self.api_endpoint}/apps/{self.app_name}", json=dict(maintenance=state)
@@ -135,13 +140,17 @@ def migrate(self):
         print(self.ini_file)
 
         if self.needs_migrate():
-            self.set_maintenance(True)
+            was_in_maintenance = self.get_maintenance()
+            if not was_in_maintenance:
+                self.set_maintenance(True)
             self.scale_down()
             sleep(30)
             self.alembic()
             self.scale_up()
             sleep(30)
-            self.set_maintenance(False)
+            if not was_in_maintenance:
+                # Don't disable maintenance mode if it was enabled externally.
+                self.set_maintenance(False)
         else:
             print("Database migration is not needed")
 

pyramid_heroku/tests/test_migrate.py

@@ -2,6 +2,7 @@
 
 from mock import call
 
+import json
 import mock
 import pytest
 import responses
@@ -84,6 +85,17 @@ def test_scale_up(self, out):
         h.scale_up()
         out.assert_has_calls([call("Scaled up to:"), call("web=5")])
 
+    @responses.activate
+    def test_get_maintenance(self):
+        h = self.Heroku("test", "etc/production.ini")
+        responses.add(  # noqa
+            responses.GET,
+            "https://api.heroku.com/apps/test",
+            status=200,
+            body=json.dumps({"maintenance": True}),
+        )
+        assert h.get_maintenance()
+
     @mock.patch("pyramid_heroku.migrate.print")
     @responses.activate
     def test_maintenance_on(self, out):
@@ -206,9 +218,42 @@ def test_migrate_skip(self, sleep, out, sub):
     @mock.patch("pyramid_heroku.migrate.sleep")
     @mock.patch("pyramid_heroku.migrate.Session")
     @responses.activate
-    def test_migrate(self, ses, sleep, out, sub):
+    def test_migrate_with_maintenance_mode(self, ses, sleep, out, sub):
+
+        h = self.Heroku("test", "etc/production.ini")
+        h.get_maintenance = mock.Mock(return_value=False)
+        h.set_maintenance = mock.Mock()
+        h.migrate()
+
+        h.set_maintenance.assert_has_calls([call(True), call(False)])
+        sub.run.assert_called_with(
+            [
+                "alembic",
+                "-c",
+                "etc/alembic.ini",
+                "-x",
+                "ini=etc/production.ini",
+                "upgrade",
+                "head",
+            ],
+            stdout=mock.ANY,
+            stderr=mock.ANY,
+        )
+
+    @mock.patch("pyramid_heroku.migrate.Heroku.set_maintenance")
+    @mock.patch("pyramid_heroku.migrate.subprocess")
+    @mock.patch("pyramid_heroku.migrate.print")
+    @mock.patch("pyramid_heroku.migrate.sleep")
+    @mock.patch("pyramid_heroku.migrate.Session")
+    @responses.activate
+    def test_migrate_skip_setting_maintenance_mode(self, ses, sleep, out, sub, set_maintenance):
+
         h = self.Heroku("test", "etc/production.ini")
+        h.get_maintenance = mock.Mock(return_value=True)
+        h.set_maintenance = mock.Mock()
         h.migrate()
+
+        assert not set_maintenance.called
         sub.run.assert_called_with(
             [
                 "alembic",
@@ -223,6 +268,7 @@ def test_migrate(self, ses, sleep, out, sub):
             stderr=mock.ANY,
         )
 
+
     @mock.patch("pyramid_heroku.migrate.subprocess")
     @mock.patch("pyramid_heroku.migrate.print")
     @mock.patch("pyramid_heroku.migrate.sleep")