How can I get setUpTestData to tear down / rollback primary keys (postgreSQL)?

Question

I am using postgreSQL 12.

I have two TestCase classes:

class TestA(TestCase):

    @classmethod
    def setUpTestData(cls):
        for _ in range(5):
            MyModel.objects.create()

    def test_1(self):
        print('Start test 1')
        objects = MyModel.objects.all()
        for i in objects:
            print(i.pk)
        self.assertEqual(True, True)

class TestB(TestCase):

    @classmethod
    def setUpTestData(cls):
        for _ in range(5):
            MyModel.objects.create()

    def test_2(self):
        print('Start test 2')
        objects = MyModel.objects.all()
        for i in objects:
            print(i.pk)
        self.assertEqual(True, True)

This results in printing:

Start test 1
1 2 3 4 5

Start test 2
6 7 8 9 10

When these tests are run it is clear that although the tearDownClass is removing the objects, the autoincrement on primary key is not being reset.

Indeed, if I try to set MyModel.objects.create(pk=1) on test_2 then I get a duplicate key value error.

This causes problems if you use something like factory_boy to create your fixtures, as it will throw similar duplicate key errors.

Why are primary keys not removed and how to I fix this situation?

PKs are not removed because Postgres PKs are generated from Sequence objects that never roll back - this makes handling concurrent inserts easier. You can [reset the sequence using ALTER SEQUENCE](https://stackoverflow.com/questions/4678110/how-to-reset-sequence-in-postgres-and-fill-id-column-with-new-data). Let me know if that solves your problem. — snakecharmerb, Jul 07 '20 at 16:35
You can reset the sequence as @snakecharmerb suggests, but your tests will be more robust if you change them so that they don't rely on hardcoded pks. — Alasdair, Jul 07 '20 at 16:36
Thanks, ok I understand the cause. However when using something like `factory_boy` in each setUpClass to generate objects it throws a duplicate key value error if the same model is used in each TestCase class; is this expected? — alias51, Jul 07 '20 at 16:38

score 0 · Answer 1 · answered Jul 08 '20 at 09:00

As pointed in comments, the reason for not resetting the PK counter is based on Postgres handling of primary keys as a sequence.

Regarding factory_boy, having been initially designed for Django (and turned into a general project later), it has built-in features to handle that: factory.Sequence.

A sequence is initialized once when the factory class is declared; each subsequent use of that factory in the same program run will reuse and increment that sequence counter.

If you wish to force the value of a unique field (either the PK or some other identifier), you can use the following:

class MyModelFactory(factory.django.DjangoModelFactory):
    class Meta:
        model = MyModel

    pk = factory.Sequence(lambda n: n)
    identifier = factory.Sequence(lambda n: 'ID-{}'.format(n))

Under the hood, some advanced logic is included to get a reasonable experience when inheriting from factories for different models; for details, see https://factoryboy.readthedocs.io/en/latest/reference.html#sharing.

However, forcing PKs in tests is seldom useful — in production, they will be managed by the database engine itself; if an external, well-known ID is required, I'd rather put it into a separate, non-auto-increment, column.

Finally, if you need a guaranteed sequential numbering scheme (e.g for accounting), you'll have to depart from PostgreSQL sequences — see this discussion for details. A great implementation can be found in the django-sequences package

How can I get setUpTestData to tear down / rollback primary keys (postgreSQL)?

1 Answers1