Database fixture patterns in PyTest - Build an Automation Script

Test user creation and retrieval using database fixture

Preconditions (3)

Step 1: Use a pytest fixture to connect to the database and start a transaction

Step 2: Insert a user with username 'testuser' and email '[e-mail]' into the 'users' table

Step 3: Query the 'users' table to retrieve the inserted user by username

Step 4: Verify the retrieved user's email matches '[e-mail]'

Step 5: Rollback the transaction to leave the database unchanged

✅ Expected Result: The inserted user is found with the correct email, and the database remains unchanged after the test

Automation Requirements - pytest

Assertions Needed:

Assert that the retrieved user's email equals '[e-mail]'

Best Practices:

Use pytest fixtures to manage database connection and transaction lifecycle

Use explicit commit or rollback to keep database clean after tests

Use parameterized queries to avoid SQL injection

Keep test data isolated and cleaned up after test

Automated Solution

PyTest

import pytest
import psycopg2
from psycopg2.extras import RealDictCursor

@pytest.fixture(scope='function')
def db_connection():
    conn = psycopg2.connect(
        dbname='testdb', user='testuser', password='testpass', host='localhost'
    )
    conn.autocommit = False
    cursor = conn.cursor(cursor_factory=RealDictCursor)
    yield cursor
    conn.rollback()
    cursor.close()
    conn.close()

def test_insert_and_retrieve_user(db_connection):
    cursor = db_connection
    insert_query = """
        INSERT INTO users (username, email) VALUES (%s, %s)
    """
    cursor.execute(insert_query, ('testuser', '[e-mail]'))

    select_query = """
        SELECT username, email FROM users WHERE username = %s
    """
    cursor.execute(select_query, ('testuser',))
    user = cursor.fetchone()

    assert user is not None, "User should be found in database"
    assert user['email'] == '[e-mail]', f"Expected email '[e-mail]', got {user['email']}"

The db_connection fixture creates a database connection and cursor with a transaction started. It yields the cursor to the test, then rolls back any changes to keep the database clean.

The test test_insert_and_retrieve_user uses this fixture to insert a user with parameterized queries to avoid SQL injection. Then it queries the user back and asserts the email matches the expected value.

Using a fixture for connection management and rollback ensures tests do not affect each other or the real data, following best practices for database testing.

Common Mistakes - 4 Pitfalls

Not rolling back or cleaning up database changes after test

Hardcoding SQL queries with string concatenation

Opening and closing database connections inside each test

Using global or module-level fixtures with shared state

Bonus Challenge

Now add data-driven testing with 3 different user inputs for username and email

Show Hint

Practice

1. What is the main purpose of using database fixtures in pytest?

easy

A. To speed up the database server

B. To write SQL queries inside test functions

C. To prepare and clean test data automatically before and after tests

D. To replace the need for assertions in tests

Solution

1. Step 1: Understand what fixtures do

   Fixtures in pytest are used to set up and tear down resources needed for tests, such as database data.

2. Step 2: Identify the role of database fixtures

   Database fixtures specifically prepare test data before tests run and clean it up after tests finish, ensuring tests run reliably.

3. Final Answer:

   To prepare and clean test data automatically before and after tests -> Option C

4. Quick Check:

   Database fixtures = setup and cleanup [OK]

Hint: Fixtures handle setup and cleanup automatically [OK]

Common Mistakes:

• Thinking fixtures run SQL queries inside tests
• Believing fixtures speed up the database server
• Confusing fixtures with assertions

2. Which of the following is the correct way to write a pytest fixture that sets up a database connection and tears it down after the test using yield?

easy

A. def db():
conn = connect()
yield conn
conn.close()

B. def db():
conn = connect()
conn.close()
yield conn

C. def db():
yield connect()
conn.close()

D. def db():
conn = connect()
return conn
conn.close()

Solution

1. Step 1: Understand yield usage in fixtures

   Using yield in a fixture splits setup (before yield) and teardown (after yield).

2. Step 2: Check each option's order

   def db(): conn = connect() yield conn conn.close() sets up connection, yields it, then closes connection after test. Others close before yield or have unreachable code.

3. Final Answer:

   def db():\n conn = connect()\n yield conn\n conn.close() -> Option A

4. Quick Check:

   Setup before yield, teardown after yield [OK]

Hint: Yield separates setup and teardown in fixtures [OK]

Common Mistakes:

• Closing connection before yield
• Placing code after return (unreachable)
• Yielding before setup

3. Given the following pytest fixture and test, what will be printed when the test runs?

import pytest

@pytest.fixture
def sample_db():
    data = {'count': 0}
    yield data
    data['count'] += 1


def test_increment(sample_db):
    print(sample_db['count'])
    sample_db['count'] += 5
    print(sample_db['count'])

medium

A. 1\n6

B. 0\n5

C. 0\n0

D. 5\n10

Solution

1. Step 1: Analyze fixture setup and teardown

   The fixture yields data with 'count' 0. After test, it increments 'count' by 1 (not affecting test output).

2. Step 2: Trace test function prints

   First print shows initial 0. Then 'count' is increased by 5, so second print shows 5.

3. Final Answer:

   0\n5 -> Option B

4. Quick Check:

   Yielded data count = 0, incremented in test = 5 [OK]

Hint: Yield returns setup data; teardown runs after test [OK]

Common Mistakes:

• Thinking teardown runs before test prints
• Assuming fixture modifies data before yield
• Confusing fixture teardown with test code

4. Identify the error in this pytest fixture that is supposed to setup a test database and clean it after tests:

@pytest.fixture
def test_db():
    conn = connect_db()
    conn.execute('CREATE TABLE users')
    return conn
    conn.execute('DROP TABLE users')
    conn.close()

medium

A. The cleanup code after return is never executed

B. The fixture should use yield instead of return for cleanup

C. The table creation SQL is incorrect

D. The fixture is missing the @pytest.mark decorator

Solution

1. Step 1: Check the fixture structure

   Code after return statement is unreachable and will never run.

2. Step 2: Understand cleanup execution

   Cleanup code must run after test, so it should be placed after yield or before return, but not after return.

3. Final Answer:

   The cleanup code after return is never executed -> Option A

4. Quick Check:

   Code after return is unreachable [OK]

Hint: Code after return in fixture won't run [OK]

Common Mistakes:

• Thinking return allows cleanup after it
• Confusing yield and return usage
• Ignoring unreachable code warnings

5. You want to create a pytest fixture that sets up a test database with multiple tables and ensures all tables are dropped after tests, even if a test fails. Which pattern best achieves this?

hard

A. Create tables once globally without cleanup to speed up tests

B. Create tables inside each test and drop them at the end of each test

C. Use return in fixture to return connection, then drop tables in a separate teardown function

D. Use a fixture with yield: create tables before yield, drop tables after yield

Solution

1. Step 1: Understand reliable setup and teardown

   Using yield in fixtures allows setup before tests and guaranteed cleanup after, even if tests fail.

2. Step 2: Evaluate options for cleanup guarantee

   Use a fixture with yield: create tables before yield, drop tables after yield uses yield to create tables before tests and drop them after, ensuring cleanup always runs.

3. Final Answer:

   Use a fixture with yield: create tables before yield, drop tables after yield -> Option D

4. Quick Check:

   Yield fixture ensures setup and guaranteed teardown [OK]

Hint: Yield fixtures guarantee cleanup after tests [OK]

Common Mistakes:

• Skipping cleanup causing leftover tables
• Relying on test code for cleanup
• Avoiding yield and missing teardown