Skip to content

Commit e056b91

Browse files
committed
test: Fix most tests
1 parent 45c0dc2 commit e056b91

File tree

3 files changed

+68
-65
lines changed

3 files changed

+68
-65
lines changed

src/sentry/spans/consumers/process_segments/factory.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -55,8 +55,7 @@ def _process_message(message: Message[KafkaPayload]):
5555
def explode_segment(spans: list[dict[str, Any]]):
5656
for span in spans:
5757
if span is not None:
58-
payload = rapidjson.dumps(span)
59-
yield Value(KafkaPayload(None, payload, []), {}, None)
58+
yield Value(KafkaPayload(None, rapidjson.dumps(span), []), {}, None)
6059

6160

6261
class DetectPerformanceIssuesStrategyFactory(ProcessingStrategyFactory[KafkaPayload]):
@@ -78,9 +77,7 @@ def __init__(
7877
cluster_name = get_topic_definition(Topic.SNUBA_SPANS)["cluster"]
7978
producer_config = get_kafka_producer_cluster_options(cluster_name)
8079
self.producer = KafkaProducer(build_kafka_configuration(default_config=producer_config))
81-
self.output_topic = ArroyoTopic(
82-
get_topic_definition(Topic.BUFFERED_SEGMENTS)["real_topic_name"]
83-
)
80+
self.output_topic = ArroyoTopic(get_topic_definition(Topic.SNUBA_SPANS)["real_topic_name"])
8481

8582
def create_with_partitions(
8683
self,

tests/sentry/spans/consumers/process_segments/test_factory.py

Lines changed: 46 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,18 @@
44
from arroyo.backends.kafka import KafkaPayload
55
from arroyo.types import BrokerValue, Message, Partition
66
from arroyo.types import Topic as ArroyoTopic
7+
from sentry_kafka_schemas.codecs import Codec
8+
from sentry_kafka_schemas.schema_types.snuba_spans_v1 import SpanEvent
79

8-
from sentry.conf.types.kafka_definition import Topic
10+
from sentry.conf.types.kafka_definition import Topic, get_topic_codec
911
from sentry.spans.consumers.process_segments.factory import DetectPerformanceIssuesStrategyFactory
1012
from sentry.testutils.helpers.options import override_options
1113
from sentry.utils import json
1214
from sentry.utils.kafka_config import get_topic_definition
1315
from tests.sentry.spans.consumers.process.test_factory import build_mock_span
1416

17+
SNUBA_SPANS_CODEC: Codec[SpanEvent] = get_topic_codec(Topic.SNUBA_SPANS)
18+
1519

1620
def build_mock_message(data, topic=None):
1721
message = mock.Mock()
@@ -26,58 +30,66 @@ def build_mock_message(data, topic=None):
2630
"standalone-spans.process-segments-consumer.enable": True,
2731
}
2832
)
29-
@mock.patch("sentry.spans.consumers.process_segments.factory.process_segment")
33+
@mock.patch(
34+
"sentry.spans.consumers.process_segments.factory.process_segment", side_effect=lambda x: x
35+
)
3036
def test_segment_deserialized_correctly(mock_process_segment):
3137
topic = ArroyoTopic(get_topic_definition(Topic.BUFFERED_SEGMENTS)["real_topic_name"])
3238
partition_1 = Partition(topic, 0)
3339
partition_2 = Partition(topic, 1)
3440
mock_commit = mock.Mock()
35-
strategy = DetectPerformanceIssuesStrategyFactory(
41+
factory = DetectPerformanceIssuesStrategyFactory(
3642
num_processes=2,
3743
input_block_size=1,
3844
max_batch_size=2,
3945
max_batch_time=1,
4046
output_block_size=1,
41-
).create_with_partitions(
42-
commit=mock_commit,
43-
partitions={},
4447
)
4548

46-
span_data = build_mock_span(project_id=1, is_segment=True)
47-
segment_data = {"spans": [span_data]}
48-
message = build_mock_message(segment_data, topic)
49+
with mock.patch.object(factory, "producer", new=mock.Mock()) as mock_producer:
50+
strategy = factory.create_with_partitions(
51+
commit=mock_commit,
52+
partitions={},
53+
)
54+
55+
span_data = build_mock_span(project_id=1, is_segment=True)
56+
segment_data = {"spans": [span_data]}
57+
message = build_mock_message(segment_data, topic)
4958

50-
strategy.submit(
51-
Message(
52-
BrokerValue(
53-
KafkaPayload(b"key", message.value().encode("utf-8"), []),
54-
partition_1,
55-
1,
56-
datetime.now(),
59+
strategy.submit(
60+
Message(
61+
BrokerValue(
62+
KafkaPayload(b"key", message.value().encode("utf-8"), []),
63+
partition_1,
64+
1,
65+
datetime.now(),
66+
)
5767
)
5868
)
59-
)
6069

61-
strategy.submit(
62-
Message(
63-
BrokerValue(
64-
KafkaPayload(b"key", message.value().encode("utf-8"), []),
65-
partition_2,
66-
1,
67-
datetime.now(),
70+
strategy.submit(
71+
Message(
72+
BrokerValue(
73+
KafkaPayload(b"key", message.value().encode("utf-8"), []),
74+
partition_2,
75+
1,
76+
datetime.now(),
77+
)
6878
)
6979
)
70-
)
7180

72-
calls = [
73-
mock.call({partition_1: 2}),
74-
mock.call({partition_2: 2}),
75-
]
81+
calls = [
82+
mock.call({partition_1: 2}),
83+
mock.call({partition_2: 2}),
84+
]
7685

77-
mock_commit.assert_has_calls(calls=calls, any_order=True)
86+
strategy.poll()
87+
strategy.join(1)
88+
strategy.terminate()
7889

79-
strategy.poll()
80-
strategy.join(1)
81-
strategy.terminate()
90+
mock_commit.assert_has_calls(calls=calls, any_order=True)
91+
assert mock_process_segment.call_args.args[0] == segment_data["spans"]
8292

83-
assert mock_process_segment.call_args.args[0] == segment_data["spans"]
93+
assert mock_producer.produce.call_count == 2
94+
assert mock_producer.produce.call_args.args[0] == ArroyoTopic("snuba-spans")
95+
SNUBA_SPANS_CODEC.decode(mock_producer.produce.call_args.args[1].value)

tests/sentry/spans/consumers/process_segments/test_message.py

Lines changed: 20 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -50,30 +50,14 @@ def repeating_span():
5050

5151
return spans
5252

53-
@override_options(
54-
{
55-
"standalone-spans.detect-performance-problems.enable": True,
56-
}
57-
)
58-
def test_n_plus_one_issue_detection(self):
59-
spans = self.generate_n_plus_one_spans()
60-
job = process_segment(spans)[0]
61-
62-
assert (
63-
job["performance_problems"][0].fingerprint
64-
== "1-GroupType.PERFORMANCE_N_PLUS_ONE_DB_QUERIES-f906d576ffde8f005fd741f7b9c8a35062361e67-1019"
65-
)
66-
67-
assert job["performance_problems"][0].type == PerformanceStreamedSpansGroupTypeExperimental
68-
6953
@override_options(
7054
{
7155
"standalone-spans.detect-performance-problems.enable": True,
7256
"standalone-spans.send-occurrence-to-platform.enable": True,
7357
}
7458
)
7559
@mock.patch("sentry.issues.ingest.send_issue_occurrence_to_eventstream")
76-
def test_sends_occurrence_to_platform(self, mock_eventstream):
60+
def test_n_plus_one_issue_detection(self, mock_eventstream):
7761
spans = self.generate_n_plus_one_spans()
7862
with mock.patch(
7963
"sentry.issues.grouptype.PerformanceStreamedSpansGroupTypeExperimental.released"
@@ -82,18 +66,23 @@ def test_sends_occurrence_to_platform(self, mock_eventstream):
8266
process_segment(spans)[0]
8367

8468
mock_eventstream.assert_called_once()
85-
assert mock_eventstream.call_args[0][1].fingerprint == [
69+
70+
performance_problem = mock_eventstream.call_args[0][1]
71+
assert performance_problem.fingerprint == [
8672
md5(
8773
b"1-GroupType.PERFORMANCE_N_PLUS_ONE_DB_QUERIES-f906d576ffde8f005fd741f7b9c8a35062361e67-1019"
8874
).hexdigest()
8975
]
76+
assert performance_problem.type == PerformanceStreamedSpansGroupTypeExperimental
9077

9178
@override_options(
9279
{
9380
"standalone-spans.detect-performance-problems.enable": True,
81+
"standalone-spans.send-occurrence-to-platform.enable": True,
9482
}
9583
)
96-
def test_n_plus_one_issue_detection_without_segment_span(self):
84+
@mock.patch("sentry.issues.ingest.send_issue_occurrence_to_eventstream")
85+
def test_n_plus_one_issue_detection_without_segment_span(self, mock_eventstream):
9786
segment_span = build_mock_span(project_id=self.project.id, is_segment=False)
9887
child_span = build_mock_span(
9988
project_id=self.project.id,
@@ -128,11 +117,16 @@ def repeating_span():
128117
repeating_spans = [repeating_span() for _ in range(7)]
129118
spans = [segment_span, child_span, cause_span] + repeating_spans
130119

131-
job = process_segment(spans)[0]
132-
133-
assert (
134-
job["performance_problems"][0].fingerprint
135-
== "1-GroupType.PERFORMANCE_N_PLUS_ONE_DB_QUERIES-f906d576ffde8f005fd741f7b9c8a35062361e67-1019"
136-
)
120+
with mock.patch(
121+
"sentry.issues.grouptype.PerformanceStreamedSpansGroupTypeExperimental.released"
122+
) as mock_released:
123+
mock_released.return_value = True
124+
process_segment(spans)[0]
137125

138-
assert job["performance_problems"][0].type == PerformanceStreamedSpansGroupTypeExperimental
126+
performance_problem = mock_eventstream.call_args[0][1]
127+
assert performance_problem.fingerprint == [
128+
md5(
129+
b"1-GroupType.PERFORMANCE_N_PLUS_ONE_DB_QUERIES-f906d576ffde8f005fd741f7b9c8a35062361e67-1019"
130+
).hexdigest()
131+
]
132+
assert performance_problem.type == PerformanceStreamedSpansGroupTypeExperimental

0 commit comments

Comments
 (0)