Coverage for sm/runner.py: 83%

1from django.test.runner import DiscoverRunner 

2 

3 

4class SmTestRunner(DiscoverRunner): 

5 """ 

6 Test runner that clears stray PostgreSQL sessions before the test database 

7 is dropped. Live-server browser tests can leave per-thread sessions behind 

8 (e.g. when a request aborts mid-view and its handler thread exits without 

9 closing its connection), which otherwise makes ``DROP DATABASE`` fail with 

10 ``psycopg.errors.ObjectInUse``. 

11 """ 

12 

13 def teardown_databases(self, old_config, verbosity=1, keepdb=False): 

14 from django.db import connections 

15 

16 for alias in connections: 

17 connection = connections[alias] 

18 # NOTE: the vendor string is "postgresql" (not "postgres"). 

19 if connection.vendor != "postgresql": 

20 continue 

21 db_name = connection.settings_dict["NAME"] 

22 # Connect to the maintenance database (not the test DB). 

23 with connection._nodb_cursor() as cursor: 

24 cursor.execute( 

25 """ 

26 SELECT pid, state, left(coalesce(query, ''), 80) 

27 FROM pg_stat_activity 

28 WHERE datname = %s AND pid <> pg_backend_pid() 

29 """, 

30 [db_name], 

31 ) 

32 stragglers = cursor.fetchall() 

33 if stragglers and verbosity >= 2: 

34 for pid_, state, query in stragglers: 

35 print( 

36 f"[SmTestRunner] terminating stray session " 

37 f"{pid_} ({state}): {query}" 

38 ) 

39 with connection._nodb_cursor() as cursor: 

40 cursor.execute( 

41 """ 

42 SELECT pg_terminate_backend(pid) 

43 FROM pg_stat_activity 

44 WHERE datname = %s AND pid <> pg_backend_pid() 

45 """, 

46 [db_name], 

47 ) 

48 return super().teardown_databases( 

49 old_config, verbosity=verbosity, keepdb=keepdb 

50 )