refactoring migrations backward command

Author: dantownsend

piccolo/apps/migrations/commands/backwards.py

@@ -3,7 +3,10 @@
 import sys
 
 from piccolo.apps.migrations.auto import MigrationManager
-from piccolo.apps.migrations.commands.base import BaseMigrationManager
+from piccolo.apps.migrations.commands.base import (
+    BaseMigrationManager,
+    MigrationResult,
+)
 from piccolo.apps.migrations.tables import Migration
 
 
@@ -21,12 +24,12 @@ def __init__(
         self.clean = clean
         super().__init__()
 
-    async def run(self):
+    async def run(self) -> MigrationResult:
         await self.create_migration_table()
 
         app_modules = self.get_app_modules()
 
-        migration_modules = []
+        migration_modules = {}
 
         for app_module in app_modules:
             app_config = getattr(app_module, "APP_CONFIG")
@@ -42,8 +45,9 @@ async def run(self):
         if len(ran_migration_ids) == 0:
             # Make sure a status of 0 is returned, as we don't want this
             # to appear as an error in automated scripts.
-            print("No migrations to reverse!")
-            sys.exit(0)
+            message = "No migrations to reverse!"
+            print(message)
+            return MigrationResult(success=True, message=message)
 
         #######################################################################
 
@@ -55,10 +59,12 @@ async def run(self):
             earliest_migration_id = self.migration_id
 
         if earliest_migration_id not in ran_migration_ids:
-            sys.exit(
+            message = (
                 "Unrecognized migration name - must be one of "
                 f"{ran_migration_ids}"
             )
+            print(message, file=sys.stderr)
+            return MigrationResult(success=False, message=message)
 
         #######################################################################
 
@@ -99,33 +105,20 @@ async def run(self):
                 if self.clean:
                     os.unlink(migration_module.__file__)
 
+            return MigrationResult(success=True)
+
         else:  # pragma: no cover
-            sys.exit("Not proceeding.")
+            message = "Not proceeding."
+            print(message, file=sys.stderr)
+            return MigrationResult(success=False, message=message)
 
 
-async def backwards(
+async def run_backwards(
     app_name: str,
     migration_id: str = "1",
     auto_agree: bool = False,
     clean: bool = False,
-):
-    """
-    Undo migrations up to a specific migration.
-
-    :param app_name:
-        The app to reverse migrations for. Specify a value of 'all' to reverse
-        migrations for all apps.
-    :param migration_id:
-        Migrations will be reversed up to and including this migration_id.
-        Specify a value of 'all' to undo all of the migrations. Specify a
-        value of '1' to undo the most recent migration.
-    :param auto_agree:
-        If true, automatically agree to any input prompts.
-    :param clean:
-        If true, the migration files which have been run backwards are deleted
-        from the disk after completing.
-
-    """
+) -> MigrationResult:
     if app_name == "all":
         sorted_app_names = BaseMigrationManager().get_sorted_app_names()
         sorted_app_names.reverse()
@@ -149,11 +142,48 @@ async def backwards(
                     auto_agree=auto_agree,
                 )
                 await manager.run()
+            return MigrationResult(success=True)
+        else:
+            return MigrationResult(success=False, message="User cancelled")
     else:
         manager = BackwardsMigrationManager(
             app_name=app_name,
             migration_id=migration_id,
             auto_agree=auto_agree,
             clean=clean,
         )
-        await manager.run()
+        return await manager.run()
+
+
+async def backwards(
+    app_name: str,
+    migration_id: str = "1",
+    auto_agree: bool = False,
+    clean: bool = False,
+):
+    """
+    Undo migrations up to a specific migration.
+
+    :param app_name:
+        The app to reverse migrations for. Specify a value of 'all' to reverse
+        migrations for all apps.
+    :param migration_id:
+        Migrations will be reversed up to and including this migration_id.
+        Specify a value of 'all' to undo all of the migrations. Specify a
+        value of '1' to undo the most recent migration.
+    :param auto_agree:
+        If true, automatically agree to any input prompts.
+    :param clean:
+        If true, the migration files which have been run backwards are deleted
+        from the disk after completing.
+
+    """
+    response = await run_backwards(
+        app_name=app_name,
+        migration_id=migration_id,
+        auto_agree=auto_agree,
+        clean=clean,
+    )
+
+    if not response.success:
+        sys.exit(1)

tests/apps/migrations/commands/test_forwards_backwards.py

@@ -88,7 +88,7 @@ def test_forwards_backwards_single_migration(self):
                 self.assertTrue(not table_class.table_exists().run_sync())
 
     @patch("piccolo.apps.migrations.commands.forwards.print")
-    def test_forwards_unknown_migration(self, print_):
+    def test_forwards_unknown_migration(self, print_: MagicMock):
         """
         Test running an unknown migrations forwards.
         """
@@ -104,13 +104,14 @@ def test_forwards_unknown_migration(self, print_):
             in print_.mock_calls
         )
 
-    def test_backwards_unknown_migration(self):
+    @patch("piccolo.apps.migrations.commands.backwards.print")
+    def test_backwards_unknown_migration(self, print_: MagicMock):
         """
         Test running an unknown migrations backwards.
         """
         run_sync(forwards(app_name="example_app", migration_id="all"))
 
-        with self.assertRaises(SystemExit) as manager:
+        with self.assertRaises(SystemExit):
             run_sync(
                 backwards(
                     app_name="example_app",
@@ -120,29 +121,24 @@ def test_backwards_unknown_migration(self):
             )
 
         self.assertTrue(
-            manager.exception.__str__().startswith(
-                "Unrecognized migration name - must be one of "
-            )
+            print_.mock_calls[0]
+            .args[0]
+            .startswith("Unrecognized migration name - must be one of ")
         )
 
     @patch("piccolo.apps.migrations.commands.backwards.print")
-    def test_backwards_no_migrations(self, print_):
+    def test_backwards_no_migrations(self, print_: MagicMock):
         """
         Test running migrations backwards if none have been run previously.
         """
-        with self.assertRaises(SystemExit) as manager:
-            run_sync(
-                backwards(
-                    app_name="example_app",
-                    migration_id="2020-12-17T18:44:30",
-                    auto_agree=True,
-                )
+        run_sync(
+            backwards(
+                app_name="example_app",
+                migration_id="2020-12-17T18:44:30",
+                auto_agree=True,
             )
-
-        self.assertEqual(manager.exception.code, 0)
-        self.assertTrue(
-            print_.mock_calls[-1] == call("No migrations to reverse!")
         )
+        self.assertTrue(call("No migrations to reverse!") in print_.mock_calls)
 
     @patch("piccolo.apps.migrations.commands.forwards.print")
     def test_forwards_no_migrations(self, print_: MagicMock):