Skip to content

Commit f4318b4

Browse files
john-z-yangjohn-z-yang
authored
test(metrics): add more features to send_metrics.py (#69545)
### Overview Adds a few more features to the send_metrics.py script. - Allow b64 encoding of sets and distributions - Allows user to specify which metrics type they would like to send
1 parent 8e37705 commit f4318b4

File tree

1 file changed

+139
-79
lines changed

1 file changed

+139
-79
lines changed

‎bin/send_metrics.py

Lines changed: 139 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -1,84 +1,121 @@
11
# pylint: skip-file
22
# flake8: noqa
33

4+
import base64
45
import datetime
6+
import functools
57
import itertools
68
import json
79
import pprint
810
import random
911
import string
12+
import struct
1013

1114
import click
1215
from arroyo.backends.kafka import KafkaPayload, KafkaProducer
1316
from arroyo.types import Topic
1417

1518
from sentry.sentry_metrics.use_case_id_registry import UseCaseID
1619

17-
make_counter_payload = lambda use_case, org_id, rand_str: {
18-
"name": f"c:{use_case}/{use_case}@none",
19-
"tags": {
20-
"environment": "production",
21-
"session.status": "init",
22-
f"metric_e2e_{use_case}_counter_k_{rand_str}": f"metric_e2e_{use_case}_counter_v_{rand_str}",
23-
},
24-
"timestamp": int(datetime.datetime.now(tz=datetime.UTC).timestamp()),
25-
"type": "c",
26-
"value": 1,
27-
"org_id": org_id,
28-
"retention_days": 90,
29-
"project_id": 3,
30-
}
31-
32-
make_dist_payload = lambda use_case, org_id, rand_str, value_len: {
33-
"name": f"d:{use_case}/duration@second",
34-
"tags": {
35-
"environment": "production",
36-
"session.status": "healthy",
37-
f"metric_e2e_{use_case}_dist_k_{rand_str}": f"metric_e2e_{use_case}_dist_v_{rand_str}",
38-
},
39-
"timestamp": int(datetime.datetime.now(tz=datetime.UTC).timestamp()),
40-
"type": "d",
41-
"value": {"format": "array", "data": [random.random() for _ in range(value_len)]},
42-
"org_id": org_id,
43-
"retention_days": 90,
44-
"project_id": 3,
45-
}
46-
47-
make_set_payload = lambda use_case, org_id, rand_str, value_len: {
48-
"name": f"s:{use_case}/error@none",
49-
"tags": {
50-
"environment": "production",
51-
"session.status": "errored",
52-
f"metric_e2e_{use_case}_set_k_{rand_str}": f"metric_e2e_{use_case}_set_v_{rand_str}",
53-
},
54-
"timestamp": int(datetime.datetime.now(tz=datetime.UTC).timestamp()),
55-
"type": "s",
56-
"value": {"format": "array", "data": [random.randint(0, 2048) for _ in range(value_len)]},
57-
"org_id": org_id,
58-
"retention_days": 90,
59-
"project_id": 3,
60-
}
61-
62-
make_gauge_payload = lambda use_case, org_id, rand_str: {
63-
"name": f"s:{use_case}/error@none",
64-
"tags": {
65-
"environment": "production",
66-
"session.status": "errored",
67-
f"metric_e2e_{use_case}_set_k_{rand_str}": f"metric_e2e_{use_case}_set_v_{rand_str}",
68-
},
69-
"timestamp": int(datetime.datetime.now(tz=datetime.UTC).timestamp()),
70-
"type": "g",
71-
"value": {
72-
"min": 1,
73-
"max": 1,
74-
"sum": 1,
75-
"count": 1,
76-
"last": 1,
77-
},
78-
"org_id": org_id,
79-
"retention_days": 90,
80-
"project_id": 3,
81-
}
20+
21+
def make_counter_payload(use_case, org_id, rand_str):
22+
return {
23+
"name": f"c:{use_case}/{use_case}@none",
24+
"tags": {
25+
"environment": "production",
26+
"session.status": "init",
27+
f"metric_e2e_{use_case}_counter_k_{rand_str}": f"metric_e2e_{use_case}_counter_v_{rand_str}",
28+
},
29+
"timestamp": int(datetime.datetime.now(tz=datetime.UTC).timestamp()),
30+
"type": "c",
31+
"value": 1,
32+
"org_id": org_id,
33+
"retention_days": 90,
34+
"project_id": 3,
35+
}
36+
37+
38+
def make_dist_payload(use_case, org_id, rand_str, value_len, b64_encode):
39+
nums = [random.random() for _ in range(value_len)]
40+
return {
41+
"name": f"d:{use_case}/duration@second",
42+
"tags": {
43+
"environment": "production",
44+
"session.status": "healthy",
45+
f"metric_e2e_{use_case}_dist_k_{rand_str}": f"metric_e2e_{use_case}_dist_v_{rand_str}",
46+
},
47+
"timestamp": int(datetime.datetime.now(tz=datetime.UTC).timestamp()),
48+
"type": "d",
49+
"value": (
50+
{
51+
"format": "base64",
52+
"data": base64.b64encode(struct.pack(f"<{len(nums)}d", *nums)),
53+
}
54+
if b64_encode
55+
else {
56+
"format": "array",
57+
"data": nums,
58+
}
59+
),
60+
"org_id": org_id,
61+
"retention_days": 90,
62+
"project_id": 3,
63+
}
64+
65+
66+
def make_set_payload(use_case, org_id, rand_str, value_len, b64_encode):
67+
INT_WIDTH = 4
68+
nums = [random.randint(0, 2048) for _ in range(value_len)]
69+
return {
70+
"name": f"s:{use_case}/error@none",
71+
"tags": {
72+
"environment": "production",
73+
"session.status": "errored",
74+
f"metric_e2e_{use_case}_set_k_{rand_str}": f"metric_e2e_{use_case}_set_v_{rand_str}",
75+
},
76+
"timestamp": int(datetime.datetime.now(tz=datetime.UTC).timestamp()),
77+
"type": "s",
78+
"value": (
79+
{
80+
"format": "base64",
81+
"data": base64.b64encode(
82+
b"".join([num.to_bytes(INT_WIDTH, byteorder="little") for num in nums])
83+
),
84+
}
85+
if b64_encode
86+
else {
87+
"format": "array",
88+
"data": nums,
89+
}
90+
),
91+
"org_id": org_id,
92+
"retention_days": 90,
93+
"project_id": 3,
94+
}
95+
96+
97+
def make_gauge_payload(use_case, org_id, rand_str):
98+
return {
99+
"name": f"s:{use_case}/error@none",
100+
"tags": {
101+
"environment": "production",
102+
"session.status": "errored",
103+
f"metric_e2e_{use_case}_set_k_{rand_str}": f"metric_e2e_{use_case}_set_v_{rand_str}",
104+
},
105+
"timestamp": int(datetime.datetime.now(tz=datetime.UTC).timestamp()),
106+
"type": "g",
107+
"value": {
108+
"min": 1,
109+
"max": 1,
110+
"sum": 1,
111+
"count": 1,
112+
"last": 1,
113+
},
114+
"org_id": org_id,
115+
"retention_days": 90,
116+
"project_id": 3,
117+
}
118+
82119

83120
make_psql = (
84121
lambda rand_str, is_generic: f"""
@@ -144,6 +181,9 @@ def produce_msgs(messages, is_generic, host, dryrun, quiet):
144181

145182

146183
@click.command()
184+
@click.option(
185+
"--metric-types", default="cdsg", show_default=True, help="The types of metrics to send"
186+
)
147187
@click.option(
148188
"--use-cases",
149189
multiple=True,
@@ -197,8 +237,24 @@ def produce_msgs(messages, is_generic, host, dryrun, quiet):
197237
show_default=True,
198238
help="Number of elements for metrics (sets and distributions).",
199239
)
240+
@click.option(
241+
"--b64-encode",
242+
default=True,
243+
show_default=True,
244+
help="Encode sets and distribution metrics values in base64",
245+
)
200246
def main(
201-
use_cases, rand_str, host, dryrun, quiet, start_org_id, end_org_id, num_bad_msg, value_len
247+
metric_types,
248+
use_cases,
249+
rand_str,
250+
host,
251+
dryrun,
252+
quiet,
253+
start_org_id,
254+
end_org_id,
255+
num_bad_msg,
256+
value_len,
257+
b64_encode,
202258
):
203259
if UseCaseID.SESSIONS.value in use_cases and len(use_cases) > 1:
204260
click.secho(
@@ -208,32 +264,36 @@ def main(
208264
)
209265
exit(1)
210266

211-
rand_str = rand_str or "".join(random.choices(string.ascii_uppercase + string.digits, k=8))
212-
213267
is_generic = UseCaseID.SESSIONS.value not in use_cases
268+
metric_types = "".join(set(metric_types))
269+
rand_str = rand_str or "".join(random.choices(string.ascii_uppercase + string.digits, k=8))
270+
payload_generators = {
271+
"c": functools.partial(make_counter_payload, rand_str=rand_str),
272+
"d": functools.partial(
273+
make_dist_payload, rand_str=rand_str, value_len=value_len, b64_encode=b64_encode
274+
),
275+
"s": functools.partial(
276+
make_set_payload, rand_str=rand_str, value_len=value_len, b64_encode=b64_encode
277+
),
278+
"g": functools.partial(make_gauge_payload, rand_str=rand_str),
279+
}
214280

215281
messages = list(
216282
itertools.chain.from_iterable(
217283
(
218-
make_counter_payload(use_case, org, rand_str),
219-
make_dist_payload(use_case, org, rand_str, value_len),
220-
make_set_payload(use_case, org, rand_str, value_len),
221-
make_gauge_payload(use_case, org, rand_str),
284+
payload_generators[metric_type](use_case=use_case, org_id=org_id)
285+
for metric_type in metric_types
222286
)
223287
for use_case in use_cases
224-
for org in range(start_org_id, end_org_id + 1)
288+
for org_id in range(start_org_id, end_org_id + 1)
225289
)
226290
)
227-
228291
messages.extend([{"BAD_VALUE": rand_str, "idx": i} for i in range(num_bad_msg)])
229292

230293
random.shuffle(messages)
231-
232294
produce_msgs(messages, is_generic, host, dryrun, quiet)
233295

234-
metrics_per_use_case = 4
235296
strs_per_use_case = 3
236-
237297
print(
238298
f"Use the following SQL to verify postgres, "
239299
f"there should be {strs_per_use_case} strings for each use cases, "
@@ -244,8 +304,8 @@ def main(
244304
if is_generic:
245305
print(
246306
f"Use the following SQL to verify clickhouse, "
247-
f"there should be {metrics_per_use_case} metrics for each use cases, "
248-
f"{metrics_per_use_case * len(use_cases) * (end_org_id - start_org_id + 1)} in total."
307+
f"there should be {len(metric_types)} metrics for each use cases, "
308+
f"{len(metric_types) * len(use_cases) * (end_org_id - start_org_id + 1)} in total."
249309
)
250310
print(make_csql(rand_str, is_generic))
251311

0 commit comments

Comments
 (0)