From ca94909cccf7bc28f7e76055ef5a3bd926b6b17b Mon Sep 17 00:00:00 2001 From: Matthias Date: Sun, 15 Feb 2026 18:09:05 +0100 Subject: [PATCH] test: improve test for migrate_set_sequence_id's --- tests/persistence/test_migrations.py | 30 +++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/tests/persistence/test_migrations.py b/tests/persistence/test_migrations.py index 31b58068c..1e99060c7 100644 --- a/tests/persistence/test_migrations.py +++ b/tests/persistence/test_migrations.py @@ -374,11 +374,39 @@ def test_migrate_get_last_sequence_ids(): def test_migrate_set_sequence_ids(): engine = MagicMock() - engine.begin = MagicMock() + # make engine.begin() usable as a context manager that returns `conn` + conn = MagicMock() + engine.begin = MagicMock( + return_value=MagicMock( + __enter__=MagicMock(return_value=conn), + ) + ) engine.name = "postgresql" + set_sequence_ids(engine, 22, 55, 5, 3, 1) + # begin called once and connection.execute invoked for each provided sequence id assert engine.begin.call_count == 1 + assert conn.execute.call_count == 5 + assert ( + conn.execute.call_args_list[0][0][0].text == "ALTER SEQUENCE orders_id_seq RESTART WITH 22" + ) + assert ( + conn.execute.call_args_list[1][0][0].text == "ALTER SEQUENCE trades_id_seq RESTART WITH 55" + ) + assert ( + conn.execute.call_args_list[2][0][0].text + == "ALTER SEQUENCE pairlocks_id_seq RESTART WITH 5" + ) + assert ( + conn.execute.call_args_list[3][0][0].text + == 'ALTER SEQUENCE "KeyValueStore_id_seq" RESTART WITH 3' + ) + assert ( + conn.execute.call_args_list[4][0][0].text + == "ALTER SEQUENCE trade_custom_data_id_seq RESTART WITH 1" + ) + engine.reset_mock() engine.begin.reset_mock()