Coverage for /Users/buh/.pyenv/versions/3.12.9/envs/es-testbed/lib/python3.12/site-packages/es_testbed/presets/searchable_test/functions.py: 69%
29 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-04-17 22:23 -0600
« prev ^ index » next coverage.py v7.6.12, created at 2025-04-17 22:23 -0600
1"""Module that will generate docs for Searchable Snapshot Test"""
3import typing as t
4import random
5import string
6import logging
7from datetime import datetime, timezone
9logger = logging.getLogger(__name__)
12def iso8601_now() -> str:
13 """
14 :returns: An ISO8601 timestamp based on now
15 :rtype: str
17 Because Python 3.12 now requires non-naive timezone declarations, we must change.
19 The new way:
21 .. code-block:: python
23 datetime.now(timezone.utc).isoformat()
25 Result: ``2024-04-16T16:00:00+00:00``
27 Note that the +00:00 is appended now where we affirmatively declare the
28 UTC timezone
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)
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 """
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
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)))
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()
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 }