This article is a follow-on from an article I wrote back in September 2017 about unit testing Postgres. In that article, I looked at how we went about testing our database schemas using the pgTAP testing framework.

In this article, I want to focus on the approach we have taken to testing functions in Postgres using the same framework. Here we will build upon the knowledge from the original article and explore some of the many functions available in pgTAP and how we can employ them to test a procedure. Before you begin, I suggest you read the original article to get pgTAP set up on your database if you have not already done so as I will assume that this is present and configured correctly.

Functions are a very powerful feature in Postgres that can be written in a variety of languages including PL/pgSQL , c , python , and others. They can be used to efficiently work with data held in the database by performing complex operations in close proximity to the data without the overheads of moving that data around. At Gradient, we are big advocates of functions and use them widely.

Below is an example function that is used to accept an invite to our Gurn application.

create function accept_invite_to_list(userId integer, userVersion integer, inviteId integer)

returns boolean

language plpgsql

as $$

declare

invite lists_invites%rowtype;

begin

if not is_valid_user(userId, userVersion) then

raise exception '30001-USER VERSION IS NOT VALID';

end if;



select * into invite from lists_invites where id = inviteId;



if invite is null then

raise exception '30010-TOKEN IS INVALID';

end if;



if invite.account_id != userId then

raise exception '30002-USER DOES NOT HAVE ACCESS';

end if;



if invite.invite_expires < now() then

raise exception '30006-TOKEN EXPIRED';

end if;



-- Grant permission, add the list to their schedule, and remove the invite.

insert into list_permissions (list_id, account_id, permission_group_id)

values (invite.list_id, userId, invite.permission_group_id);



perform add_list_to_default_context(userId, invite.list_id);



delete from lists_invites where id = inviteId;



return true;

end;

$$;

Let’s quickly walk through this function and take a look at what it is doing. As you can see, the function initially performs some checks to ensure it is able to proceed, raising exceptions where appropriate. These checks enable us to ensure that the user is in a valid state, the invite exists, the user accepting the invite is the invited user, and the invite is still valid. Once we have performed these checks successfully we proceed to process the invite, setup permissions and return a true value.

We need to test this function end to end from its signature right through to ensuring that the success state and each of the error states is invoked correctly and the expected result is returned. I want to take you through how we can test each aspect of this function, in turn, to give us the confidence that it is behaving as intended.

1. Function signature

create function accept_invite_to_list(userId integer, userVersion integer, inviteId integer)

returns boolean

language plpgsql

as $$

Focusing in on the function signature, we can see that the function has a name accept_invite_to_list it also has three integer arguments userId , userVersion , and inviteId . The next line in our function outlines the return type of the function, as you can see we are returning a boolean value. Finally, we have the language that the function is written in PL/pgSQL . To help us test a function signature pgTAP provides some useful utility functions

has_function()

This function checks whether the named function exists in the database with the defined signature.

function_lang_is()

This function checks whether a named function with the given signature is written using the specified language.

function_returns()

This function checks whether a named function returns the specified type.

We can combine these pgTAP utility functions to test that our function definition is accurate and as expected. This forms the first block of our suite of tests.

-- #################################################

-- 1. Function definition checks

-- #################################################

SELECT has_function(

'accept_invite_to_list',

ARRAY ['integer', 'integer', 'integer'],

'accept_invite_to_list exists'

); SELECT function_lang_is(

'accept_invite_to_list',

ARRAY [ 'integer', 'integer', 'integer'], 'plpgsql',

'accept_invite_to_list is plpgsql'

); SELECT function_returns(

'accept_invite_to_list',

'boolean'

);

2. Seed data

Before we can go on to test our function is working as expected, we need to seed some data. The second step of our test suite is where we populate this test data. Our tests will be run within a transaction so anything we add in this phase will not persist beyond the scope of the test. If you have common shared data then you may wish to set this up once in advance of the tests. I personally prefer starting each test with a clean schema and being explicit when adding and deleting data for my tests to remove the risk of any possible side effects.

Below is the PL/pgSQL we use to seed the database for this suite of tests. As you can see, this step simply involves inserting data before checking that the data is present in the database using the following pgTAP utility function.

is()

This function checks whether the result of a particular column value matches an expected value.

The other thing to note here is the use of set which allows us to declare values within the scope of the transaction. These values are then referenced using the :variable_name format, enabling them to be queried within SQL. We can couple this with Postgres’s ::type suffix notation to enable us to type cast correctly.

The combination of these techniques allows us to seed data and check that it is available with the correct types as outlined below.

-- #################################################

-- 2. Seed data for function tests

-- #################################################

\set account_id 5

\set invite_id 10

\set invite_error_id 100

\set list_id 99



insert into accounts

(id, username, email, validated)

values (:account_id, 'test', 'test@gurn.io', 1);



insert into lists

(id, list_name, description, created_by, updated_by, owned_by)

values (:list_id::integer, 'test'::text, 'a test list'::text, 1::integer, 1::integer, 1::integer);



SELECT is(

setup_user_list(:account_id::integer),

true,

'Setup default lists'

);



SELECT is(

accounts.id,

:account_id::integer,

'Seed an account'

) from accounts where id = :account_id;



insert into lists_invites

(id, account_id, list_id, permission_group_id, created_by, invite_expires)

values (:invite_id::integer, :account_id::integer, :list_id::integer, 3::integer, 1::integer, now() + interval '1 day');



SELECT is(

lists_invites.id,

:invite_id::integer,

'Seed an invite'

) from lists_invites where id = :invite_id;

3. Test our success cases

After we have seeded data for testing, the third component of our test is proving our success cases. Here we are using three pgTAP utility functions

ok()

This function checks that a supplied SQL statement executes without error.

results_eq()

This function checks that a supplied SQL statement results set matches a given example.

is_empty()

This function checks that the result of a sql statement is empty. This is useful for checking records have been deleted from tables.

You’ll notice that we are using PREPARE statements here to construct SQL statements with interpolated values. A PREPARE statement generates a SQL string that can be reused, we are having to use them here as results_eq() only accepts a SQL string value and we cannot interpolate our variable values into a string in this way.

Our success check invokes our function and checks the state of the database after the invocation to ensure that everything has behaved as expected.

-- #################################################

-- 3. Test function invocation success cases

-- #################################################

-- Checks the success cases and expected database

-- state after procedure execution SELECT ok(

accept_invite_to_list(:account_id::integer, 0::integer, :invite_id::integer ),

'Invoke accept_invite_to_list'

);



-- Output checks

-- Use prepared statements to support interpolation

PREPARE list_permission_checks AS select list_id, account_id, permission_group_id, version from list_permissions where account_id = :account_id::integer and list_id = :list_id::integer; PREPARE list_permission_results AS select :list_id::integer, :account_id::integer, 3::integer, 0::integer;



SELECT results_eq(

'list_permission_checks',

'list_permission_results',

'list_permissions should be set for user'

);



SELECT is_empty(

'select * from lists_invites',

'lists_invites should be deleted'

);

4. Test our error cases

We now need to test our various error cases to ensure that they are being thrown correctly. We do this using the following pgTAP utility function.

throws_ok()

This function checks that an error is thrown and it matches the provided string and/or error code.

In order to test some of our error states, we need to generate extra seed data. We do this using a combination of insert statements and PREPARE statements as described in our seed data section.

Our first two tests can be invoked by simply passing incorrect values directly to the function. In the first case, we supply a deliberately erroneous userVersion number that raises an exception. In the second case, we provide an erroneous inviteToken that also raises an error.

In order to test user access and expired tokens, we need to create a new invite for a different user account and insert it into the database. Additionally, to test that the invite expired error is raised correctly we set the expiration date to one day before the test is being conducted.

SELECT throws_ok(

'select accept_invite_to_list(1::integer, 10::integer, 100::integer )',

'30001-USER VERSION IS NOT VALID',

'Throws user version not valid exception'

);



SELECT throws_ok(

'select accept_invite_to_list(1::integer, 0::integer, 100::integer )',

'30010-TOKEN IS INVALID',

'Throws token is invalid exception'

);



insert into lists_invites

(id, account_id, list_id, permission_group_id, created_by, invite_expires)

values (:invite_error_id::integer, :account_id::integer, :list_id::integer, 3::integer, 1::integer, now() - interval '1 day');



PREPARE list_invite_no_access AS select accept_invite_to_list(1::integer, 0::integer, :invite_error_id::integer );

SELECT throws_ok(

'list_invite_no_access',

'30002-USER DOES NOT HAVE ACCESS',

'Throws user does not have access exception'

);



PREPARE list_invite_expired AS select accept_invite_to_list(:account_id::integer, 0::integer, :invite_error_id::integer );

SELECT throws_ok(

'list_invite_expired',

'30006-TOKEN EXPIRED',

'Throws token expired exception'

);

5. Wrap it all up

As with any pgTAP test, it is good practice to clean up any artefacts of the test that may exist beyond the scope of the transaction. In our case, the PREPARE statements that we defined and bound to the connection and are not destroyed after a test is run. This means that if you need to re-run a test on the same database connection you will likely run into errors about a prepare statement already existing. An example error is shown below.

[42P05] ERROR: prepared statement "list_permission_checks" already exists

In order to prevent this from occurring, we use the DEALLOCATE command to tell Postgres to delete the prepared statement as part of our clean up process.

-- #################################################

-- 5. Finish tests and clean up

-- #################################################

SELECT *

FROM finish();



deallocate list_permission_results;

deallocate list_permission_checks;

deallocate list_invite_no_access;

deallocate list_invite_expired;

Pulling it all together

In this article we have looked at how we can use some of the functions exposed by pgTAP to test the various aspects of a function including its signature, return type, success cases and error cases. We have also explored how we can effectively seed data and ensure that this process worked as expected.

Bringing these steps together, our final unit test for out accept_list_invite procedure is below. As with all pgTAP tests you need to specify a plan stating the number of tests that you are executing. You may also notice the use of the BEGIN and ROLLBACK statements, these simply start and revert the database transaction.

We can invoke our test suite by running it directly within our database or by using the pg_prove command line utility.

-- Start transaction and plan the tests.

BEGIN;



-- Plan count should be the number of tests

SELECT plan(13);



-- #################################################

-- 1. Function definition checks

-- #################################################

SELECT has_function(

'accept_invite_to_list',

ARRAY ['integer', 'integer', 'integer'],

'accept_invite_to_list exists'

); SELECT function_lang_is(

'accept_invite_to_list',

ARRAY [ 'integer', 'integer', 'integer'], 'plpgsql',

'accept_invite_to_list is plpgsql'

); SELECT function_returns(

'accept_invite_to_list',

'boolean'

);



-- #################################################

-- 2. Seed data for function tests

-- #################################################

\set local account_id 5

\set local invite_id 10

\set local invite_error_id 100

\set local list_id 99



insert into accounts

(id, username, email, validated)

values (:account_id, 'test', 'test@gurn.io', 1);



insert into lists

(id, list_name, description, created_by, updated_by, owned_by)

values (:list_id::integer, 'test'::text, 'a test list'::text, 1::integer, 1::integer, 1::integer);



SELECT is(

setup_user_list(:account_id::integer),

true,

'Setup default lists'

);



SELECT is(

accounts.id,

:account_id::integer,

'Seed an account'

) from accounts where id = :account_id;



insert into lists_invites

(id, account_id, list_id, permission_group_id, created_by, invite_expires)

values (:invite_id::integer, :account_id::integer, :list_id::integer, 3::integer, 1::integer, now() + interval '1 day');



SELECT is(

lists_invites.id,

:invite_id::integer,

'Seed an invite'

) from lists_invites where id = :invite_id;



-- #################################################

-- 3. Test function invocation success cases

-- #################################################

-- Checks the success cases and expected database

-- state after procedure execution

SELECT ok(

accept_invite_to_list(:account_id::integer, 0::integer, :invite_id::integer ),

'Invoke accept_invite_to_list'

);



-- Output checks

-- Use prepared statements to support interpolation

PREPARE list_permission_checks AS select list_id, account_id, permission_group_id, version from list_permissions where account_id = :account_id::integer and list_id = :list_id::integer; PREPARE list_permission_results AS select :list_id::integer, :account_id::integer, 3::integer, 0::integer;

-- For just hard-coded results could use $$VALUES (1, 'test') etc but as we are interpolating we can't



SELECT results_eq(

'list_permission_checks',

'list_permission_results',

'list_permissions should be set for user'

);



SELECT is_empty(

'select * from lists_invites',

'lists_invites should be deleted'

);



-- #################################################

-- 4. Check function invocation error cases

-- #################################################

SELECT throws_ok(

'select accept_invite_to_list(1::integer, 10::integer, 100::integer )',

'30001-USER VERSION IS NOT VALID',

'Throws user version not valid exception'

);



SELECT throws_ok(

'select accept_invite_to_list(1::integer, 0::integer, 100::integer )',

'30010-TOKEN IS INVALID',

'Throws token is invalid exception'

);



insert into lists_invites

(id, account_id, list_id, permission_group_id, created_by, invite_expires)

values (:invite_error_id::integer, :account_id::integer, :list_id::integer, 3::integer, 1::integer, now() - interval '1 day');



PREPARE list_invite_no_access AS select accept_invite_to_list(1::integer, 0::integer, :invite_error_id::integer ); SELECT throws_ok(

'list_invite_no_access',

'30002-USER DOES NOT HAVE ACCESS',

'Throws user does not have access exception'

);



PREPARE list_invite_expired AS select accept_invite_to_list(:account_id::integer, 0::integer, :invite_error_id::integer ); SELECT throws_ok(

'list_invite_expired',

'30006-TOKEN EXPIRED',

'Throws token expired exception'

);



-- #################################################

-- 5. Finish tests and clean up

-- #################################################

SELECT *

FROM finish();



deallocate list_permission_results;

deallocate list_permission_checks;

deallocate list_invite_no_access;

deallocate list_invite_expired;



ROLLBACK;

Adding tests to functions does not have to be done all at once. You can slowly introduce them over time. It will give you confidence in the performance of critical functions and help ensure no regressions are introduced when iterating on your application or working to introduce new features.

We hope you find this useful. We are always interested to hear how others approach this problem. If you have any questions, please leave them in the comments below and we will get back to you.