From 2be534c24771f8ba3fdb8ec88fda3f314d57ae3b Mon Sep 17 00:00:00 2001
From: Nathan Louie <nxlouie@umich.edu>
Date: Mon, 17 Oct 2022 21:38:50 -0700
Subject: [PATCH] add tests

---
 alembic/command.py       | 10 ++++++++--
 alembic/util/__init__.py |  2 +-
 alembic/util/exc.py      |  3 +++
 tests/test_command.py    | 33 ++++++++++++++++++++++++++++++++-
 4 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/alembic/command.py b/alembic/command.py
index 5d86cd53..d90e998c 100644
--- a/alembic/command.py
+++ b/alembic/command.py
@@ -172,6 +172,10 @@ def revision(
      the other parameters, this option is only available via programmatic
      use of :func:`.command.revision`
 
+    :param check: instead of generating a revision, checks if this revision
+     will contain upgrade ops; no new ops will have no action; new ops will
+     error; this is the ``--check`` option to ``alembic revision``.
+
     """
 
     script_directory = ScriptDirectory.from_config(config)
@@ -238,11 +242,13 @@ def revision(
     
     if check:
         if not autogenerate:
-            util.err("check flag cannot be used without autogenerate flag.")
+            raise util.CommandError(
+                "Check flag cannot be used without autogenerate flag"
+            )
         migration_script = revision_context.generated_revisions[-1]
         diffs = migration_script.upgrade_ops.as_diffs()
         if diffs:
-            util.err(f"Revision has upgrade ops to run: {diffs}.")
+            raise util.RevisionOpsNotEmptyError(f"Revision has upgrade ops to run: {diffs}.")
         else:
             log.info("Revision has no upgrade ops to run.")      
     else:
diff --git a/alembic/util/__init__.py b/alembic/util/__init__.py
index d5fa4d32..f77184bf 100644
--- a/alembic/util/__init__.py
+++ b/alembic/util/__init__.py
@@ -1,5 +1,5 @@
 from .editor import open_in_editor
-from .exc import CommandError
+from .exc import CommandError, RevisionOpsNotEmptyError
 from .langhelpers import _with_legacy_names
 from .langhelpers import asbool
 from .langhelpers import dedupe_tuple
diff --git a/alembic/util/exc.py b/alembic/util/exc.py
index f7ad0211..bd999415 100644
--- a/alembic/util/exc.py
+++ b/alembic/util/exc.py
@@ -1,2 +1,5 @@
 class CommandError(Exception):
     pass
+
+class RevisionOpsNotEmptyError(Exception):
+    pass
\ No newline at end of file
diff --git a/tests/test_command.py b/tests/test_command.py
index 0c0ce378..761f221f 100644
--- a/tests/test_command.py
+++ b/tests/test_command.py
@@ -8,8 +8,9 @@ import re
 from typing import cast
 
 from sqlalchemy import exc as sqla_exc
-from sqlalchemy import text
+from sqlalchemy import VARCHAR, text
 from sqlalchemy.engine import Engine
+from sqlalchemy.sql.schema import Column
 
 from alembic import __version__
 from alembic import command
@@ -385,6 +386,36 @@ finally:
             self.cfg,
             autogenerate=True,
         )
+  
+    def test_rev_check_with_no_autogen(self):
+        self._env_fixture()
+        assert_raises_message(
+            util.CommandError,
+            "Check flag cannot be used without autogenerate flag",
+            command.revision,
+            self.cfg,
+            autogenerate=False,
+            check=True,
+        )
+
+    def test_rev_autogen_check_no_changes(self):
+        self._env_fixture()
+        command.revision(self.cfg, autogenerate=True, check=True) # no problem
+
+    def test_rev_autogen_check_changes_detected(self):
+        self._env_fixture()
+        with mock.patch(
+            "alembic.operations.ops.UpgradeOps.as_diffs",
+            return_value=[('remove_column', None, 'foo', Column('old_data', VARCHAR()))]
+        ):
+            assert_raises_message(
+                util.RevisionOpsNotEmptyError,
+                "Revision has upgrade ops to run:",
+                command.revision,
+                self.cfg,
+                autogenerate=True,
+                check=True,
+            )
 
     def test_pk_constraint_normally_prevents_dupe_rows(self):
         self._env_fixture()
-- 
2.47.3