Coverage for django_query_capture/test_utils.py: 100%

1""" 

2Test utility to help you check Duplicate, Similar, and Slow queries in the test 

3""" 

4import typing 

5 

6from contextlib import ContextDecorator, ExitStack 

7 

8from django.test import override_settings 

9from django.utils.module_loading import import_string 

10 

11from django_query_capture import BasePresenter, query_capture 

12from django_query_capture.settings import get_config 

13from django_query_capture.utils import CaptureStdOutToString 

14 

15 

16class AssertInefficientQuery(ContextDecorator): 

17 def __init__( 

18 self, 

19 num: typing.Optional[int] = None, 

20 seconds: typing.Optional[int] = None, 

21 ignore_patterns: typing.Optional[typing.List[str]] = None, 

22 ): 

23 """ 

24 Args: 

25 num: `Duplicate`, `Similar` Threshold, The value of the setting is ignored. 

26 seconds: `Slow` Threshold, The value of the setting is ignored. 

27 ignore_patterns: A list of patterns to ignore IGNORE_SQL_PATTERNS of settings. 

28 """ 

29 self.ignore_patterns = ignore_patterns or get_config()["IGNORE_SQL_PATTERNS"] 

30 self.num = num 

31 self.seconds = seconds 

32 self.presenter_cls: typing.Type[BasePresenter] = import_string( 

33 get_config()["PRESENTER"] 

34 ) 

35 

36 def __enter__(self): 

37 """ 

38 Run [query_capture.__enter__][decorators.query_capture] and `override_settings.__enter__` 

39 At this time, the `ignore_output=True` is set so that the output of [query_capture][decorators.query_capture] can be ignored. 

40 override_settings are used to ignore existing set values within the current context. 

41 """ 

42 self._exit_stack = ExitStack().__enter__() 

43 self.query_capture = query_capture(ignore_output=True) 

44 config = get_config().copy() 

45 config.update( 

46 { 

47 "PRINT_THRESHOLDS": { 

48 **config["PRINT_THRESHOLDS"], 

49 "SLOW_MIN_SECOND": self.seconds, 

50 "DUPLICATE_MIN_COUNT": self.num, 

51 "SIMILAR_MIN_COUNT": self.num, 

52 } 

53 } 

54 ) 

55 self._exit_stack.enter_context(override_settings(QUERY_CAPTURE=config)) 

56 self._exit_stack.enter_context(self.query_capture) 

57 return self.query_capture 

58 

59 def __exit__(self, exc_type, exc_val, exc_tb): 

60 """ 

61 End the context of [query_capture][decorators.query_capture] and override_settings. 

62 And if there is an item above the threshold, the test fails and the failed content is printed. 

63 """ 

64 self._exit_stack.close() 

65 

66 with CaptureStdOutToString() as stdout: 

67 self.presenter_cls(self.query_capture.classifier).print() 

68 result = stdout.getvalue() 

69 if self.query_capture.classifier["has_over_threshold"]: 

70 raise AssertionError(result)