Simplify migration tests

tests/conftest.py

@@ -1,6 +1,7 @@
+import contextlib
+import functools
 import os
 import signal as stdlib_signal
-from contextlib import closing
 
 import aiopg
 import psycopg2
@@ -18,61 +19,78 @@
         os.environ.pop(key)
 
 
-def _execute(cursor, query, *identifiers):
-    cursor.execute(
-        sql.SQL(query).format(
+def cursor_execute(cursor, query, *identifiers, format=True):
+    if identifiers:
+        query = sql.SQL(query).format(
             *(sql.Identifier(identifier) for identifier in identifiers)
         )
-    )
+    cursor.execute(query)
+
+
+@contextlib.contextmanager
+def db_executor(dbname):
+    with contextlib.closing(psycopg2.connect("", dbname=dbname)) as connection:
+        connection.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
+        with connection.cursor() as cursor:
+            yield functools.partial(cursor_execute, cursor)
+
+
+@pytest.fixture
+def db_execute():
+    return db_executor
+
+
+def db_create(dbname, template=None):
+    with db_executor("postgres") as execute:
+        execute("DROP DATABASE IF EXISTS {}", dbname)
+        if template:
+            execute("CREATE DATABASE {} TEMPLATE {}", dbname, template)
+        else:
+            execute("CREATE DATABASE {}", dbname)
+
+
+def db_drop(dbname):
+    with db_executor("postgres") as execute:
+        execute("DROP DATABASE IF EXISTS {}", dbname)
+
+
+@pytest.fixture
+def db_factory():
+    dbs_to_drop = []
+
+    def _(dbname, template=None):
+        db_create(dbname=dbname, template=template)
+        dbs_to_drop.append(dbname)
+
+    yield _
+
+    for dbname in dbs_to_drop:
+        db_drop(dbname=dbname)
 
 
 @pytest.fixture(scope="session")
 def setup_db():
 
-    with closing(psycopg2.connect("", dbname="postgres")) as connection:
-        connection.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
-        with connection.cursor() as cursor:
-            _execute(
-                cursor, "DROP DATABASE IF EXISTS {}", "procrastinate_test_template"
-            )
-            _execute(cursor, "CREATE DATABASE {}", "procrastinate_test_template")
+    dbname = "procrastinate_test_template"
+    db_create(dbname=dbname)
 
-    connector = aiopg_connector.PostgresConnector(dbname="procrastinate_test_template")
+    connector = aiopg_connector.PostgresConnector(dbname=dbname)
     schema_manager = schema.SchemaManager(connector=connector)
     schema_manager.apply_schema()
     # We need to close the psycopg2 underlying connection synchronously
     connector.close()
 
-    with closing(
-        psycopg2.connect("", dbname="procrastinate_test_template")
-    ) as connection:
-        connection.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
-        yield connection
+    yield dbname
 
-    with closing(psycopg2.connect("", dbname="postgres")) as connection:
-        connection.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
-        with connection.cursor() as cursor:
-            _execute(
-                cursor, "DROP DATABASE IF EXISTS {}", "procrastinate_test_template"
-            )
+    db_drop(dbname=dbname)
 
 
 @pytest.fixture
-def connection_params(setup_db):
-    with setup_db.cursor() as cursor:
-        _execute(cursor, "DROP DATABASE IF EXISTS {}", "procrastinate_test")
-        _execute(
-            cursor,
-            "CREATE DATABASE {} TEMPLATE {}",
-            "procrastinate_test",
-            "procrastinate_test_template",
-        )
+def connection_params(setup_db, db_factory):
+    db_factory(dbname="procrastinate_test", template=setup_db)
 
     yield {"dsn": "", "dbname": "procrastinate_test"}
 
-    with setup_db.cursor() as cursor:
-        _execute(cursor, "DROP DATABASE IF EXISTS {}", "procrastinate_test")
-
 
 @pytest.fixture
 async def connection(connection_params):

tests/migration/pgservice.ini

@@ -0,0 +1,4 @@
+[procrastinate_schema]
+dbname=procrastinate_schema
+[procrastinate_migrations]
+dbname=procrastinate_migrations

tests/migration/test_migration.py

@@ -1,114 +1,77 @@
 import os
 import subprocess
-from contextlib import closing
 
-import pkg_resources
-import psycopg2
 import pytest
-from psycopg2 import sql
-from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
 
-from procrastinate import __version__, aiopg_connector, schema
+from procrastinate import aiopg_connector, schema
 
-PG_SERVICES = """
-[procrastinate_test1]
-dbname=procrastinate_test1
-[procrastinate_test2]
-dbname=procrastinate_test2
-"""
-
-
-def _execute_sql(cursor, query, *identifiers):
-    cursor.execute(
-        sql.SQL(query).format(
-            *(sql.Identifier(identifier) for identifier in identifiers)
-        )
-    )
+BASELINE = "0.5.0"
 
 
+# Pum config file is currently broken (https://github.com/opengisch/pum/issues/5)
+# When it's fixed, we can add a config file here, and then get rid of omit_table_dir
 @pytest.fixture
-def migrations_directory():
-    return pkg_resources.resource_filename("procrastinate", "sql/migrations")
+def pum():
 
+    env = {**os.environ, "PGSERVICEFILE": "tests/migration/pgservice.ini"}
 
-@pytest.fixture
-def procrastinate_version():
-    try:
-        idx = __version__.index("+")
-    except ValueError:
-        idx = None
-    return __version__[:idx]
+    def _(command, args: str, omit_table_dir=False):
+        first_args = ["pum", command]
+        if not omit_table_dir:
+            first_args.extend(
+                ["--table=public.pum", "--dir=procrastinate/sql/migrations"]
+            )
+        return subprocess.run(first_args + args.split(), check=True, env=env)
 
-
-@pytest.fixture
-def pg_service_file(tmp_path):
-    file_path = tmp_path / "pg_service.conf"
-    with open(file_path, "w") as file_:
-        file_.write(PG_SERVICES)
-    yield file_path
+    return _
 
 
 @pytest.fixture
-def environment(pg_service_file):
-    env = os.environ.copy()
-    env["PGSERVICEFILE"] = pg_service_file
-    return env
-
+def schema_database(db_factory, pum):
+    dbname = "procrastinate_schema"
+    db_factory(dbname=dbname)
 
-@pytest.fixture
-def setup_databases(procrastinate_version, migrations_directory, environment):
-
-    # create the procrastinate_test1 and procrastinate_test2 databases
-    with closing(psycopg2.connect("", dbname="postgres")) as connection:
-        connection.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
-        with connection.cursor() as cursor:
-            _execute_sql(cursor, "DROP DATABASE IF EXISTS {}", "procrastinate_test1")
-            _execute_sql(cursor, "CREATE DATABASE {}", "procrastinate_test1")
-            _execute_sql(cursor, "DROP DATABASE IF EXISTS {}", "procrastinate_test2")
-            _execute_sql(cursor, "CREATE DATABASE {}", "procrastinate_test2")
-
-    # apply the procrastinate schema to procrastinate_test1
-    connector = aiopg_connector.PostgresConnector(dbname="procrastinate_test1")
+    # apply the current procrastinate schema to procrastinate_schema
+    connector = aiopg_connector.PostgresConnector(dbname=dbname)
     schema_manager = schema.SchemaManager(connector=connector)
     schema_manager.apply_schema()
     connector.close()
 
-    # set the baseline version in procrastinate_test1
-    cmd = (
-        "pum baseline -p procrastinate_test1 -t public.pum "
-        f"-d {migrations_directory} -b {procrastinate_version}"
-    )
-    subprocess.run(cmd.split(), check=True, env=environment)
+    # set the baseline version in procrastinate_schema
+    # This db is as far as can be.
+    pum("baseline", f"--pg_service {dbname} --baseline 999.999.999")
 
-    # apply the baseline schema to procrastinate_test2
-    cmd = f"psql -d procrastinate_test2 -f {migrations_directory}/baseline-0.5.0.sql"
-    subprocess.run(cmd.split(), check=True)
+    return dbname
 
-    # set the baseline version in procrastinate_test2
-    cmd = (
-        f"pum baseline -p procrastinate_test2 -t public.pum -d {migrations_directory} "
-        "-b 0.5.0"
-    )
-    subprocess.run(cmd.split(), check=True, env=environment)
 
-    yield
+@pytest.fixture
+def migrations_database(db_factory, db_execute, pum):
+    dbname = "procrastinate_migrations"
+    db_factory(dbname=dbname)
 
-    with closing(psycopg2.connect("", dbname="postgres")) as connection:
-        connection.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
-        with connection.cursor() as cursor:
-            _execute_sql(cursor, "DROP DATABASE IF EXISTS {}", "procrastinate_test2")
-            _execute_sql(cursor, "DROP DATABASE IF EXISTS {}", "procrastinate_test1")
+    # apply the baseline schema to procrastinate_migrations
+    with db_execute(dbname) as execute:
+        with open(f"procrastinate/sql/migrations/baseline-{BASELINE}.sql") as file:
+            execute(file.read())
 
+    # set the baseline version in procrastinate_migrations
+    pum("baseline", f"--pg_service {dbname} --baseline {BASELINE}")
 
-def test_migration(setup_databases, migrations_directory, environment):
-    # apply the migrations on the procrastinate_test2 database
-    cmd = f"pum upgrade -p procrastinate_test2 -t public.pum -d {migrations_directory}"
-    subprocess.run(cmd.split(), check=True, env=environment)
+    return dbname
 
-    # check that the databases procrastinate_test1 and procrastinate_test2 have
+
+def test_migration(schema_database, migrations_database, pum):
+    # apply the migrations on the migrations_database database
+    pum("upgrade", f"--pg_service {migrations_database}")
+
+    # check that the schema_database and migrations_database have
     # the same schema
-    cmd = "pum check -p1 procrastinate_test1 -p2 procrastinate_test2 -v 2"
-    proc = subprocess.run(cmd.split(), env=environment, stdout=True, stderr=True)
+    proc = pum(
+        "check",
+        f"--pg_service1={schema_database} --pg_service2={migrations_database} "
+        "--verbose_level=2",
+        omit_table_dir=True,
+    )
 
     # pum check exits with a non-zero return code if the databases don't have
     # the same schema