Merge pull request #328 from jcpunk/fixture-load-schema

Add hook to centralize schema loads

README.rst

@@ -95,6 +95,18 @@ Sample test
         postgresql.commit()
         cur.close()
 
+If you want the database fixture to be automatically populated with your schema:
+
+.. code-block:: python
+
+    postgresql_my_with_schema = factories.postgresql('postgresql_my_proc', load=['schemafile.sql', 'otherschema.sql'])
+
+.. note::
+
+    The database will still be dropped each time.
+
+
+
 Connecting to already existing postgresql database
 --------------------------------------------------
 
@@ -183,6 +195,12 @@ You can pick which you prefer, but remember that these settings are handled in t
      - postgresql_dbname
      - -
      - test
+   * - Default Schema
+     - load
+     - --postgresql-load
+     - postgresql_load
+     -
+     -
    * - PostgreSQL connection options
      - options
      - --postgresql-options

src/pytest_postgresql/factories.py

@@ -36,7 +36,7 @@ def get_config(request):
     config = {}
     options = [
         'exec', 'host', 'port', 'user', 'password', 'options', 'startparams',
-        'logsprefix', 'unixsocketdir', 'dbname'
+        'logsprefix', 'unixsocketdir', 'dbname', 'load'
     ]
     for option in options:
         option_name = 'postgresql_' + option
@@ -202,12 +202,13 @@ def postgresql_noproc_fixture(request):
     return postgresql_noproc_fixture
 
 
-def postgresql(process_fixture_name, db_name=None):
+def postgresql(process_fixture_name, db_name=None, load=None):
     """
     Return connection fixture factory for PostgreSQL.
 
     :param str process_fixture_name: name of the process fixture
     :param str db_name: database name
+    :param list load: SQL to automatically load into our test database
     :rtype: func
     :returns: function which makes a connection to postgresql
     """
@@ -235,6 +236,7 @@ def postgresql_factory(request):
         pg_password = proc_fixture.password
         pg_options = proc_fixture.options
         pg_db = db_name or config['dbname']
+        pg_load = load or config['load']
 
         with DatabaseJanitor(
                 pg_user, pg_host, pg_port, pg_db, proc_fixture.version,
@@ -248,6 +250,11 @@ def postgresql_factory(request):
                 port=pg_port,
                 options=pg_options
             )
+            if pg_load:
+                for filename in pg_load:
+                    with open(filename, 'r') as _fd:
+                        with connection.cursor() as cur:
+                            cur.execute(_fd.read())
             yield connection
             connection.close()
 

src/pytest_postgresql/plugin.py

@@ -32,6 +32,7 @@
 _help_logsprefix = "Prefix for the log files"
 _help_unixsocketdir = "Location of the socket directory"
 _help_dbname = "Default database name"
+_help_load = "Load this SQL file by default, may be set multiple times"
 
 
 def pytest_addoption(parser):
@@ -96,6 +97,13 @@ def pytest_addoption(parser):
         default='tests'
     )
 
+    parser.addini(
+        name='postgresql_load',
+        type='pathlist',
+        help=_help_load,
+        default=None
+    )
+
     parser.addoption(
         '--postgresql-exec',
         action='store',
@@ -167,6 +175,13 @@ def pytest_addoption(parser):
         help=_help_dbname
     )
 
+    parser.addoption(
+        '--postgresql-load',
+        action='append',
+        dest='postgresql_load',
+        help=_help_load
+    )
+
 
 postgresql_proc = factories.postgresql_proc()
 postgresql_nooproc = factories.postgresql_noproc()

tests/conftest.py

@@ -1,8 +1,10 @@
 """Tests main conftest file."""
+import os
 from pytest_postgresql import factories
 
 
 PG_CTL = '/usr/lib/postgresql/{ver}/bin/pg_ctl'
+TEST_SQL_DIR = os.path.dirname(os.path.abspath(__file__)) + '/test_sql/'
 
 # pylint:disable=invalid-name
 postgresql92 = factories.postgresql_proc(PG_CTL.format(ver='9.2'), port=None)
@@ -15,6 +17,11 @@
 
 postgresql_proc2 = factories.postgresql_proc(port=9876)
 postgresql2 = factories.postgresql('postgresql_proc2', db_name='test-db')
+postgresql_load_1 = factories.postgresql('postgresql_proc2', db_name='test-db',
+                                         load=[TEST_SQL_DIR + 'test.sql', ])
+postgresql_load_2 = factories.postgresql('postgresql_proc2', db_name='test-db',
+                                         load=[TEST_SQL_DIR + 'test.sql',
+                                               TEST_SQL_DIR + 'test2.sql'])
 
 postgresql_rand_proc = factories.postgresql_proc(port=None)
 postgresql_rand = factories.postgresql('postgresql_rand_proc')

tests/test_postgresql.py

@@ -4,7 +4,8 @@
 import pytest
 
 
-QUERY = "CREATE TABLE test (id serial PRIMARY KEY, num integer, data varchar);"
+MAKE_Q = "CREATE TABLE test (id serial PRIMARY KEY, num integer, data varchar);"
+SELECT_Q = "SELECT * FROM test;"
 
 
 @pytest.mark.skipif(
@@ -27,24 +28,42 @@ def test_postgresql_proc(request, postgres):
 def test_main_postgres(postgresql):
     """Check main postgresql fixture."""
     cur = postgresql.cursor()
-    cur.execute(QUERY)
+    cur.execute(MAKE_Q)
     postgresql.commit()
     cur.close()
 
 
 def test_two_postgreses(postgresql, postgresql2):
     """Check two postgresql fixtures on one test."""
     cur = postgresql.cursor()
-    cur.execute(QUERY)
+    cur.execute(MAKE_Q)
     postgresql.commit()
     cur.close()
 
     cur = postgresql2.cursor()
-    cur.execute(QUERY)
+    cur.execute(MAKE_Q)
     postgresql2.commit()
     cur.close()
 
 
+def test_postgres_load_one_file(postgresql_load_1):
+    """Check postgresql fixture can load one file."""
+    cur = postgresql_load_1.cursor()
+    cur.execute(SELECT_Q)
+    results = cur.fetchall()
+    assert len(results) == 1
+    cur.close()
+
+
+def test_postgres_load_two_files(postgresql_load_2):
+    """Check postgresql fixture can load two files."""
+    cur = postgresql_load_2.cursor()
+    cur.execute(SELECT_Q)
+    results = cur.fetchall()
+    assert len(results) == 2
+    cur.close()
+
+
 def test_rand_postgres_port(postgresql_rand):
     """Check if postgres fixture can be started on random port."""
     assert postgresql_rand.status == psycopg2.extensions.STATUS_READY

tests/test_sql/test.sql

@@ -0,0 +1,2 @@
+CREATE TABLE test (id serial PRIMARY KEY, num integer, data varchar);
+INSERT INTO test VALUES(1, 2, 'c');

tests/test_sql/test2.sql

@@ -0,0 +1 @@
+INSERT INTO test VALUES(2, 1, 'z');