Coverage for /Users/buh/.pyenv/versions/3.12.2/envs/es-testbed/lib/python3.12/site-packages/es_testbed/presets/searchable_test/functions.py: 69%

29 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-08-21 12:05 -0600

1"""Module that will generate docs for Searchable Snapshot Test""" 

2 

3import typing as t 

4import random 

5import string 

6import logging 

7from datetime import datetime, timezone 

8 

9logger = logging.getLogger(__name__) 

10 

11 

12def iso8601_now() -> str: 

13 """ 

14 :returns: An ISO8601 timestamp based on now 

15 :rtype: str 

16 

17 Because Python 3.12 now requires non-naive timezone declarations, we must change. 

18 

19 The new way: 

20 

21 .. code-block:: python 

22 

23 datetime.now(timezone.utc).isoformat() 

24 

25 Result: ``2024-04-16T16:00:00+00:00`` 

26 

27 Note that the +00:00 is appended now where we affirmatively declare the 

28 UTC timezone 

29 

30 As a result, we will use this function to prune away the timezone if it is 

31 +00:00 and replace it with Z, which is shorter Zulu notation for UTC (which 

32 Elasticsearch uses) 

33 

34 We are MANUALLY, FORCEFULLY declaring timezone.utc, so it should ALWAYS be 

35 +00:00, but could in theory sometime show up as a Z, so we test for that. 

36 """ 

37 

38 parts = datetime.now(timezone.utc).isoformat().split('+') 

39 if len(parts) == 1: 

40 if parts[0][-1] == 'Z': 

41 return parts[0] # _______ It already ends with a Z for Zulu/UTC time 

42 return f'{parts[0]}Z' # _____ It doesn't end with a Z so we put one there 

43 if parts[1] == '00:00': 

44 return f'{parts[0]}Z' # _____ It doesn't end with a Z so we put one there 

45 return f'{parts[0]}+{parts[1]}' # Fallback publishes the +TZ, whatever that was 

46 

47 

48def randomstr(length: int = 16, lowercase: bool = False) -> str: 

49 """ 

50 :param length: The length of the random string 

51 :param lowercase: Whether to force the string to lowercase 

52 :returns: A random string of letters and numbers based on `length` and `lowercase` 

53 """ 

54 letters = string.ascii_uppercase 

55 if lowercase: 

56 letters = string.ascii_lowercase 

57 return str(''.join(random.choices(letters + string.digits, k=length))) 

58 

59 

60def doc_generator( 

61 count: int = 10, start_at: int = 0, match: bool = True 

62) -> t.Generator[t.Dict, None, None]: 

63 """ 

64 :param count: Create this many docs 

65 :param start_at: Append value starts with this value 

66 :param match: Do we want fieldnames to match between docgen runs, or be random? 

67 Also affects document document field "number" (value will be incremental if 

68 match is True, a random value between 1001 and 32767 if False) 

69 :returns: A generator shipping docs 

70 """ 

71 keys = ['message', 'nested', 'deep'] 

72 # Start with an empty map 

73 matchmap = {} 

74 # Iterate over each key 

75 for key in keys: 

76 # If match is True 

77 if match: 

78 # Set matchmap[key] to key 

79 matchmap[key] = key 

80 else: 

81 # Otherwise matchmap[key] will have a random string value 

82 matchmap[key] = randomstr() 

83 

84 # This is where count and start_at matter 

85 for num in range(start_at, start_at + count): 

86 yield { 

87 '@timestamp': iso8601_now(), 

88 'message': f'{matchmap["message"]}{num}', # message# or randomstr# 

89 'number': ( 

90 num if match else random.randint(1001, 32767) 

91 ), # value of num or random int 

92 'nested': {'key': f'{matchmap["nested"]}{num}'}, # nested# 

93 'deep': {'l1': {'l2': {'l3': f'{matchmap["deep"]}{num}'}}}, # deep# 

94 }