From 19f76c8c075c0371d1d961861d12a235230f8367 Mon Sep 17 00:00:00 2001
From: Arijit Basu <sayanarijit@gmail.com>
Date: Wed, 3 Feb 2021 18:00:29 +0530
Subject: [PATCH] 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                          |  8 +++++
 pyproject.toml                       |  2 +-
 pyramid_heroku/__init__.py           |  1 -
 pyramid_heroku/migrate.py            | 13 ++++++--
 pyramid_heroku/tests/test_migrate.py | 48 +++++++++++++++++++++++++++-
 5 files changed, 67 insertions(+), 5 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index deac851..2f531c4 100644
--- a/CHANGES.rst
+++ b/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
 -----
 
diff --git a/pyproject.toml b/pyproject.toml
index 5a39ee8..ae8db12 100644
--- a/pyproject.toml
+++ b/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" ]
diff --git a/pyramid_heroku/__init__.py b/pyramid_heroku/__init__.py
index ef5569c..ba7decf 100644
--- a/pyramid_heroku/__init__.py
+++ b/pyramid_heroku/__init__.py
@@ -3,7 +3,6 @@
 from ast import literal_eval
 from expandvars import expandvars
 
-import sys
 import typing as t
 
 
diff --git a/pyramid_heroku/migrate.py b/pyramid_heroku/migrate.py
index a24fb5b..95c3c23 100644
--- a/pyramid_heroku/migrate.py
+++ b/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")
 
diff --git a/pyramid_heroku/tests/test_migrate.py b/pyramid_heroku/tests/test_migrate.py
index 85832f4..cfceab8 100644
--- a/pyramid_heroku/tests/test_migrate.py
+++ b/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")