Testing Stored Procedures

Testing your database schema and migrations is crucial to ensure code behaves as expected, catch bugs early, and prevent regressions. Databases enforce logic, constraints, and complex relationships, so testing ensures these elements work correctly and remain intact after changes.

In this guide we will learn how to use Atlas's schema test command to test database stored procedures.

Stored Procedures

Stored procedures are sets of precompiled queries grouped together to perform specific tasks and are stored directly on the database.

Atlas Pro Feature

Stored procedures are currently available only to Atlas Pro users. To use this feature, run:

atlas login

Project Setup

Schema File

For this example, let's assume we have the following schema, including a stored procedure:

schema.hcl

schema "public" {
}

table "sales" {
  schema = schema.public
  column "id" {
    type = int
  }
  column "sale_amount" {
    type =  numeric
  }
  column "sale_date" {
    type = date
  }
}

table "archive_sales" {
  schema = schema.public
  column "id" {
    type = int
  }
  column "sale_amount" {
    type = int
  }
  column "sale_date" {
    type = date
  }
}

procedure "archive_old_sales" {
  schema = schema.public
  lang   = PLpgSQL
  arg "cutoff_date" {
    type = date
  }
  as = <<-SQL
  BEGIN
      -- Insert old sales into archive_sales
  INSERT INTO archive_sales (id, sale_amount, sale_date)
  SELECT id, sale_amount, sale_date
  FROM sales
  WHERE sale_date < cutoff_date;

  -- Delete old sales from sales
  DELETE FROM sales
  WHERE sale_date < cutoff_date;
  END;
  SQL
}

In the schema above we have a sales table, an archive_sales table and an archive_old_sales procedure. The procedure moves old sales from the sales table to the archive_sales table based on a cutoff date that is given when calling the procedure.

Config File

Before we begin testing, create a config file named atlas.hcl.

In this file we will create an environment, specify the source of our schema, and a URL for our dev database.

We will also create a file named schema.test.hcl to write our tests, and add it to the atlas.hcl file in the test block.

atlas.hcl

env "dev" {
  src = "file://schema.hcl"
  dev = "docker://postgres/15/dev?search_path=public"

  # Test configuration for local development.
  test {
    schema {
      src = ["schema.test.hcl"]
    }
  }
}

Writing Tests

Simple Test

Let's start off with a simple test that will:

1. Seed data into the sales table.
2. Call the archive_old_sales procedure with a cutoff date "2024-07-18".
3. Check that the archive_sales table is correctly populated.
4. Check that the archived sales no longer exist in the sales table.

schema.test.hcl

test "schema" "procedure" {
  # Seed data
  exec {
    sql = <<-SQL
      INSERT INTO sales (id, sale_amount, sale_date) VALUES
      (1, 150.00, '2024-07-18'),
      (2, 200.00, '2024-06-20'),
      (1, 350.00, '2024-07-10');
    SQL
  }
  # Execute the procedure with a specific cutoff date
  exec {
    sql = "CALL archive_old_sales('2024-07-18')"  # Archive sales before this date
  }

  # Verify data in archive_sales table
  exec {
    sql = "SELECT COUNT(*) FROM archive_sales WHERE sale_date < '2024-07-18'"
    output = "2" # Expect 2 archived sales
  }

  # Verify data in sales table
  exec {
    sql = "SELECT COUNT(*) FROM sales"
    output = "1"  # Expect 1 sale remaining in the sales table after cutoff date
  }
}

Run the test by running:

atlas schema test --env dev

The output should look similar to:

Test Output

-- PASS: procedure (4ms)
PASS

Table Driven Test

Another alternative is to write a table driven test. This test uses the for_each meta-argument, which accepts a map or a set of values and is used to generate a test case for each item in the set or map.

Following similar logic to the test above, we will execute the following:

1. Seed data into the sales table.
2. Use a for_each meta-argument to iterate over two test cases.
3. Check that the archive_sales table is correctly populated based on the cutoff date in the specific test case.
4. Check that the archived sales no longer exist in the sales table based on the cutoff date in the specific test case.

schema.test.hcl

test "schema" "procedure" {
  # Seed data
  exec {
    sql = <<-SQL
      INSERT INTO sales (id, sale_amount, sale_date) VALUES
      (1, 150.00, '2024-07-18'),
      (2, 200.00, '2024-06-20'),
      (1, 350.00, '2024-07-10');
    SQL
  }
  for_each = [
    { cutoff_date: "2024-07-18", expected_archived: "2", expected_remaining: "1" },
    { cutoff_date: "2024-08-01", expected_archived: "3", expected_remaining: "0" }
  ]

  # Execute the procedure with the current test case cutoff date
  exec {
    sql = "CALL archive_old_sales('${each.value.cutoff_date}')"
  }

  # Verify data in archive_sales table
  exec {
    sql = "SELECT COUNT(*) FROM archive_sales WHERE sale_date < '${each.value.cutoff_date}'"
    output = each.value.expected_archived
  }

  # Verify data in sales table
  exec {
    sql = "SELECT COUNT(*) FROM sales"
    output = each.value.expected_remaining
  }

  log {
    message = "Testing cutoff date:  ${each.value.cutoff_date} - Expected Archived: ${each.value.expected_archived}, Expected Remaining: ${each.value.expected_remaining}"
  }
}

Run the test by running:

atlas schema test --env dev

The output should look similar to:

Test Output

-- PASS: procedure/1 (4ms)
    schema.test.hcl:33: Testing cutoff date:  2024-07-15 - Expected Archived: 2, Expected Remaining: 1
-- PASS: procedure/2 (1ms)
    schema.test.hcl:33: Testing cutoff date:  2024-07-20 - Expected Archived: 3, Expected Remaining: 0
PASS