From 98fad1de0939dbab158ce566314030cb8ba9dede Mon Sep 17 00:00:00 2001
From: "Michael J. Sullivan" <sully@msully.net>
Date: Thu, 29 Feb 2024 22:31:50 -0800
Subject: [PATCH] Fix CI test failure by cleaning up in test_edgeql_ddl_reindex
 (#6955)

Also fix a loss of ordering in the tests caused by the `get_setup_script` change and its interaction with `get_cases_by_shard`.
---
 edb/testbase/server.py   |  4 ++--
 tests/test_edgeql_ddl.py | 50 +++++++++++++++++++++++-----------------
 2 files changed, 31 insertions(+), 23 deletions(-)

diff --git a/edb/testbase/server.py b/edb/testbase/server.py
index 7e021bf56fe..f7e170a78d7 100644
--- a/edb/testbase/server.py
+++ b/edb/testbase/server.py
@@ -2456,8 +2456,8 @@ def get_cases_by_shard(cases, selected_shard, total_shards, verbosity, stats):
     # Prepare the source heaps
     setup_count = 0
     for case, tests in cases.items():
-        setup_script = getattr(case, 'get_setup_script', lambda: None)()
-        if setup_script and tests:
+        setup_script_getter = getattr(case, 'get_setup_script', None)
+        if setup_script_getter and tests:
             tests_per_setup = []
             est_per_setup = setup_est = stats.get(
                 'setup::' + case.get_database_name(), (new_setup_est, 0),
diff --git a/tests/test_edgeql_ddl.py b/tests/test_edgeql_ddl.py
index 3cfb55948d6..a7a02c7c3f1 100644
--- a/tests/test_edgeql_ddl.py
+++ b/tests/test_edgeql_ddl.py
@@ -16199,27 +16199,35 @@ async def test_edgeql_ddl_reindex(self):
             create module test;
             create type test::Bar extending Foo;
         ''')
-        await self.con.execute('''
-            administer reindex(Foo)
-        ''')
-        await self.con.execute('''
-            administer reindex(Foo.foo)
-        ''')
-        await self.con.execute('''
-            administer reindex(Foo.bar)
-        ''')
-        await self.con.execute('''
-            administer reindex(Foo.tgt)
-        ''')
-        await self.con.execute('''
-            administer reindex(Foo.tgts)
-        ''')
-        await self.con.execute('''
-            administer reindex(test::Bar)
-        ''')
-        await self.con.execute('''
-            administer reindex(Object)
-        ''')
+        try:
+            await self.con.execute('''
+                administer reindex(Foo)
+            ''')
+            await self.con.execute('''
+                administer reindex(Foo.foo)
+            ''')
+            await self.con.execute('''
+                administer reindex(Foo.bar)
+            ''')
+            await self.con.execute('''
+                administer reindex(Foo.tgt)
+            ''')
+            await self.con.execute('''
+                administer reindex(Foo.tgts)
+            ''')
+            await self.con.execute('''
+                administer reindex(test::Bar)
+            ''')
+            await self.con.execute('''
+                administer reindex(Object)
+            ''')
+        finally:
+            await self.con.execute('''
+                drop type test::Bar;
+                drop type Foo;
+                drop type Tgt;
+                drop module test;
+            ''')
 
     async def _deadlock_tester(self, setup, teardown, modification, query):
         """Deadlock test helper.