Use --count option, refactor cli tests

peewee_migrate/cli.py

@@ -11,7 +11,6 @@
 
 VERBOSE = ['WARNING', 'INFO', 'DEBUG', 'NOTSET']
 CLEAN_RE = re.compile(r'\s+$', re.M)
-ROLLBACK_NUMBER_RE = re.compile(r'^(\d+)$')
 
 
 def get_router(directory, database, verbose=0):
@@ -86,24 +85,23 @@ def create(name, database=None, auto=False, auto_source=False, directory=None, v
 
 
 @cli.command()
-@click.argument('name', required=False, default='1')
+@click.argument('name', required=False)
+@click.option('--count', required=False, default=1, type=int)
 @click.option('--database', default=None, help="Database connection")
 @click.option('--directory', default='migrations', help="Directory where migrations are stored")
 @click.option('-v', '--verbose', count=True)
-def rollback(name, database=None, directory=None, verbose=None):
+def rollback(name, count, database=None, directory=None, verbose=None):
     """
     Rollback a migration with given name or number of migrations with given
     name as integer number
     """
     router = get_router(directory, database, verbose)
-    search = ROLLBACK_NUMBER_RE.search(name)
-    if search:
-        number_to_rollback = int(search.group(1))
-        if len(router.done) < number_to_rollback:
+    if not name:
+        if len(router.done) < count:
             raise RuntimeError(
-                'Unable to rollback %s migrations from %s: %s' % 
-                (number_to_rollback, len(router.done), router.done))
-        for _ in range(number_to_rollback):
+                'Unable to rollback %s migrations from %s: %s' %
+                (count, len(router.done), router.done))
+        for _ in range(count):
             router = get_router(directory, database, verbose)
             name = router.done[-1]
             router.rollback(name)

tests/test_cli.py

@@ -1,57 +1,100 @@
 from click.testing import CliRunner
+import pytest
+
+from peewee_migrate.cli import cli
+from peewee_migrate.cli import get_router
 
 runner = CliRunner()
 
 
-def test_cli(tmpdir):
-    tmpdir = str(tmpdir)
-    from peewee_migrate.cli import cli
-    from peewee_migrate.cli import get_router
+@pytest.fixture
+def dir_option(tmpdir):
+    return '--directory=%s' % tmpdir
+
+
+@pytest.fixture
+def db_url(tmpdir):
+    db_path = '%s/test_sqlite.db' % tmpdir
+    open(db_path, 'a').close()
+    return 'sqlite:///%s' % db_path
+
+
+@pytest.fixture
+def db_option(db_url):
+    return '--database=%s' % db_url
+
+
+@pytest.fixture
+def router(tmpdir, db_url):
+    return lambda: get_router(str(tmpdir), db_url)
 
+
+@pytest.fixture
+def migrations(router):
+    migrations_number = 5
+    name = 'test'
+    for i in range(migrations_number):
+        router().create(name)
+    return ['00%s_test' % i for i in range(1, migrations_number + 1)]
+
+
+@pytest.fixture
+def migrations_str(migrations):
+    return ', '.join(migrations)
+
+
+def test_help():
     result = runner.invoke(cli, ['--help'])
     assert result.exit_code == 0
     assert 'migrate' in result.output
     assert 'create' in result.output
     assert 'rollback' in result.output
 
-    db_path = '%s/test_sqlite.db' % tmpdir
-    db_url = 'sqlite:///%s' % db_path
-    dir_option = '--directory=%s' % tmpdir
-    db_option = '--database=%s' % db_url
 
-    open(db_path, 'a').close()
-
-    migrations_number = 5
-
-    for i in range(migrations_number):
+def test_create(dir_option, db_option):
+    for i in range(2):
         result = runner.invoke(cli, ['create', dir_option, db_option, '-vvv', 'test'])
         assert result.exit_code == 0
 
-    migrations_names = ['00%s_test' % i for i in range(1, migrations_number + 1)]
-    migrations_names_str = ', '.join(migrations_names)
-
-    # The fake seems having an issue or at least issue during testing, as
-    # it affects freshly loaded router.done and breaks the tests after it
-    # result = runner.invoke(cli, ['migrate', dir_option, db_option, '-v', '--fake'])
-    # assert result.exit_code == 0
-    # assert 'Migrations completed: %s' % migrations_names_str in result.output
-    # assert 'add_column' not in result.output
 
+def test_migrate(dir_option, db_option, migrations_str):
     result = runner.invoke(cli, ['migrate', dir_option, db_option])
     assert result.exit_code == 0
-    assert 'Migrations completed: %s' % migrations_names_str in result.output
+    assert 'Migrations completed: %s' % migrations_str in result.output
 
+
+def test_list(dir_option, db_option, migrations):
     result = runner.invoke(cli, ['list', dir_option, db_option])
-    assert 'Migrations are done:\n001_test' in result.output
+    assert 'Migrations are done:\n' in result.output
+    assert 'Migrations are undone:\n%s' % '\n'.join(migrations) in result.output
 
-    result = runner.invoke(cli, ['rollback', dir_option, db_option, '005_test'])
+
+def test_rollback(dir_option, db_option, router, migrations):
+    router().run()
+
+    result = runner.invoke(cli, ['rollback', dir_option, db_option])
     assert not result.exception
-    assert get_router(tmpdir, db_url).done == migrations_names[:-1]
+    assert router().done == migrations[:-1]
 
-    result = runner.invoke(cli, ['rollback', dir_option, db_option, '2'])
+    result = runner.invoke(cli, ['rollback', dir_option, db_option, '004_test'])
     assert not result.exception
-    assert get_router(tmpdir, db_url).done == migrations_names[:-3]
+    assert router().done == migrations[:-2]
+
+    result = runner.invoke(cli, ['rollback', dir_option, db_option, '--count=2'])
+    assert not result.exception
+    assert router().done == migrations[:-4]
 
     result = runner.invoke(cli, ['rollback', dir_option, db_option, '005_test'])
     assert result.exception
     assert result.exception.args[0] == 'Only last migration can be canceled.'
+
+
+def test_fake(dir_option, db_option, migrations_str, router):
+    result = runner.invoke(cli, ['migrate', dir_option, db_option, '-v', '--fake'])
+    assert result.exit_code == 0
+    assert 'Migrations completed: %s' % migrations_str in result.output
+
+    # TODO: Find a way of testing fake. This is unclear why the following fails.
+    # assert not router().done
+
+