Coverage for utilities/testing/query_counts.py: 0%

1import json 

2import os 

3import sys 

4import threading 

5from contextlib import contextmanager 

6from pathlib import Path 

7 

8try: 

9 import fcntl 

10except ImportError: 

11 fcntl = None 

12 

13from django.apps import apps as django_apps 

14from django.db import connection 

15from django.test.utils import CaptureQueriesContext 

16 

17__all__ = ( 

18 'assert_expected_query_count', 

19) 

20 

21UPDATE_ENV_VAR = 'UPDATE_QUERY_COUNTS' 

22BASELINE_FILENAME = 'query_counts.json' 

23 

24_loaded_baselines = {} 

25_lock = threading.Lock() 

26 

27 

28def _is_update_mode(): 

29 return bool(os.environ.get(UPDATE_ENV_VAR)) 

30 

31 

32def _is_parallel_test_run(): 

33 # Heuristic: inspects sys.argv for Django's --parallel flag. This is 

34 # sufficient for the project's standard `manage.py test` invocations but 

35 # will not detect parallelism introduced by other test runners. 

36 for arg in sys.argv: 

37 if arg == '--parallel' or arg.startswith('--parallel='): 

38 return True 

39 return False 

40 

41 

42def _baseline_path(app_label): 

43 app_config = django_apps.get_app_config(app_label) 

44 return Path(app_config.path) / 'tests' / BASELINE_FILENAME 

45 

46 

47def _load_baseline(app_label): 

48 with _lock: 

49 if app_label in _loaded_baselines: 

50 return _loaded_baselines[app_label] 

51 path = _baseline_path(app_label) 

52 if path.exists(): 

53 with path.open() as f: 

54 data = json.load(f) 

55 else: 

56 data = {} 

57 _loaded_baselines[app_label] = data 

58 return data 

59 

60 

61def _record_update(app_label, key, count): 

62 # Write the baseline file synchronously rather than buffering until process 

63 # exit, so updates are not lost if the runner terminates via os._exit() or 

64 # a signal. An OS-level exclusive lock (where available) protects against 

65 # concurrent processes — e.g. two simultaneous update-mode invocations — 

66 # clobbering one another's writes. 

67 with _lock: 

68 path = _baseline_path(app_label) 

69 path.parent.mkdir(parents=True, exist_ok=True) 

70 with path.open('a+') as f: 

71 if fcntl is not None: 

72 fcntl.flock(f.fileno(), fcntl.LOCK_EX) 

73 f.seek(0) 

74 content = f.read() 

75 existing = json.loads(content) if content else {} 

76 existing[key] = count 

77 f.seek(0) 

78 f.truncate() 

79 json.dump(existing, f, indent=2, sort_keys=True) 

80 f.write('\n') 

81 

82 

83@contextmanager 

84def assert_expected_query_count(test_case, name): 

85 """ 

86 Assert that the wrapped block performs the number of SQL queries recorded 

87 in the per-app baseline file (`<app>/tests/query_counts.json`). 

88 

89 The baseline key is `<model_label>:<name>`. By default `<model_label>` is 

90 derived from `test_case.model._meta.model_name`. Test cases that use 

91 runtime-generated models with unstable names (e.g. names derived from a 

92 database primary-key sequence) can declare a ``query_count_model_label`` 

93 class attribute to provide a stable, human-assigned label instead: 

94 

95 class MyViewTestCase(ViewTestCases.PrimaryObjectViewTestCase): 

96 query_count_model_label = 'my-stable-label' 

97 

98 When the `UPDATE_QUERY_COUNTS` environment variable is set, the assertion 

99 is skipped and the observed count is written back to the baseline file 

100 immediately. Update mode requires serial test execution (no --parallel). 

101 """ 

102 model = test_case.model 

103 app_label = model._meta.app_label 

104 label = getattr(test_case, 'query_count_model_label', None) 

105 model_name = label if label is not None else model._meta.model_name 

106 key = f'{model_name}:{name}' 

107 

108 if _is_update_mode(): 

109 if _is_parallel_test_run(): 

110 raise RuntimeError( 

111 f"{UPDATE_ENV_VAR}=1 cannot be combined with --parallel; " 

112 f"re-run serially to regenerate query-count baselines." 

113 ) 

114 ctx = CaptureQueriesContext(connection) 

115 with ctx: 

116 yield 

117 _record_update(app_label, key, len(ctx.captured_queries)) 

118 return 

119 

120 baseline = _load_baseline(app_label) 

121 if key not in baseline: 

122 test_case.fail( 

123 f"No query-count baseline recorded for {app_label}/{key}. " 

124 f"Re-run with {UPDATE_ENV_VAR}=1 (serially) to record it." 

125 ) 

126 

127 expected = baseline[key] 

128 ctx = CaptureQueriesContext(connection) 

129 with ctx: 

130 yield 

131 actual = len(ctx.captured_queries) 

132 if actual != expected: 

133 sample = '\n'.join( 

134 f" {i + 1}. {q['sql'][:240]}" 

135 for i, q in enumerate(ctx.captured_queries) 

136 ) 

137 test_case.fail( 

138 f"Query count for {app_label}/{key} changed: " 

139 f"expected {expected}, got {actual}. " 

140 f"If this change is intentional, re-run with {UPDATE_ENV_VAR}=1 " 

141 f"to update the baseline.\nObserved queries:\n{sample}" 

142 )