Skip to content

Commit b779b1d

Browse files
abdihakim92x1jrdh
authored andcommitted
Added full unit coverage for seed_random
1 parent ed278e0 commit b779b1d

File tree

1 file changed

+347
-0
lines changed

1 file changed

+347
-0
lines changed
Lines changed: 347 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,347 @@
1+
from collections.abc import Iterator
2+
from contextlib import ExitStack, nullcontext
3+
from types import SimpleNamespace
4+
from unittest import mock
5+
6+
import pytest
7+
from django.core.management import CommandParser
8+
9+
from metrics.interfaces.management.commands.seed_random import Command, SCALE_CONFIGS
10+
11+
MODULE_PATH = "metrics.interfaces.management.commands.seed_random"
12+
13+
14+
def _fake_metric_hierarchy() -> SimpleNamespace:
15+
theme = SimpleNamespace(name="Theme 1")
16+
sub_theme = SimpleNamespace(name="SubTheme 1", theme=theme)
17+
topic = SimpleNamespace(name="Topic 1", sub_theme=sub_theme)
18+
return SimpleNamespace(name="Metric 1", topic=topic)
19+
20+
21+
def _fake_geography() -> SimpleNamespace:
22+
geography_type = SimpleNamespace(name="Nation")
23+
return SimpleNamespace(
24+
name="Area 1",
25+
geography_code="RND0001",
26+
geography_type=geography_type,
27+
)
28+
29+
30+
class TestSeedRandomCommand:
31+
def test_add_arguments_parses_defaults(self):
32+
parser = CommandParser(prog="manage.py seed_random")
33+
34+
Command().add_arguments(parser)
35+
options = parser.parse_args([])
36+
37+
assert options.dataset == "both"
38+
assert options.scale == "small"
39+
assert options.seed is None
40+
assert options.truncate_first is False
41+
42+
@mock.patch(f"{MODULE_PATH}.random.seed")
43+
@mock.patch(f"{MODULE_PATH}.call_command")
44+
@mock.patch.object(Command, "_seed_metrics_data")
45+
@mock.patch.object(Command, "_print_summary")
46+
@mock.patch(f"{MODULE_PATH}.time.perf_counter")
47+
def test_handle_metrics_dataset(
48+
self,
49+
spy_perf_counter: mock.MagicMock,
50+
spy_print_summary: mock.MagicMock,
51+
spy_seed_metrics_data: mock.MagicMock,
52+
spy_call_command: mock.MagicMock,
53+
spy_random_seed: mock.MagicMock,
54+
):
55+
spy_perf_counter.side_effect = [10.0, 14.5]
56+
spy_seed_metrics_data.return_value = {
57+
"Theme": 3,
58+
"SubTheme": 6,
59+
"Topic": 12,
60+
"Metric": 10,
61+
"Geography": 5,
62+
"CoreTimeSeries": 1,
63+
"APITimeSeries": 1,
64+
}
65+
66+
Command().handle(dataset="metrics", scale="small", truncate_first=True, seed=42)
67+
68+
spy_random_seed.assert_called_once_with(42)
69+
spy_seed_metrics_data.assert_called_once_with(
70+
scale_config=SCALE_CONFIGS["small"],
71+
truncate_first=True,
72+
)
73+
spy_call_command.assert_not_called()
74+
spy_print_summary.assert_called_once_with(
75+
dataset="metrics",
76+
scale="small",
77+
seed=42,
78+
counts=spy_seed_metrics_data.return_value,
79+
runtime_seconds=4.5,
80+
)
81+
82+
@mock.patch(f"{MODULE_PATH}.random.seed")
83+
@mock.patch(f"{MODULE_PATH}.call_command")
84+
@mock.patch.object(Command, "_seed_metrics_data")
85+
@mock.patch.object(Command, "_print_summary")
86+
@mock.patch(f"{MODULE_PATH}.time.time")
87+
@mock.patch(f"{MODULE_PATH}.time.perf_counter")
88+
def test_handle_cms_dataset_uses_time_seed_and_builds_cms(
89+
self,
90+
spy_perf_counter: mock.MagicMock,
91+
spy_time: mock.MagicMock,
92+
spy_print_summary: mock.MagicMock,
93+
spy_seed_metrics_data: mock.MagicMock,
94+
spy_call_command: mock.MagicMock,
95+
spy_random_seed: mock.MagicMock,
96+
):
97+
spy_perf_counter.side_effect = [20.0, 22.0]
98+
spy_time.return_value = 1234
99+
100+
Command().handle(dataset="cms", scale="large", truncate_first=False, seed=None)
101+
102+
spy_random_seed.assert_called_once_with(1234)
103+
spy_seed_metrics_data.assert_not_called()
104+
spy_call_command.assert_called_once_with("build_cms_site")
105+
spy_print_summary.assert_called_once_with(
106+
dataset="cms",
107+
scale="large",
108+
seed=1234,
109+
counts={
110+
"Theme": 0,
111+
"SubTheme": 0,
112+
"Topic": 0,
113+
"Metric": 0,
114+
"Geography": 0,
115+
"CoreTimeSeries": 0,
116+
"APITimeSeries": 0,
117+
},
118+
runtime_seconds=2.0,
119+
)
120+
121+
@mock.patch.object(Command, "_truncate_metrics_data")
122+
@mock.patch.object(Command, "_seed_time_series_rows")
123+
@mock.patch.object(Command, "_bulk_create")
124+
@mock.patch(f"{MODULE_PATH}.Geography")
125+
@mock.patch(f"{MODULE_PATH}.Metric")
126+
@mock.patch(f"{MODULE_PATH}.Topic")
127+
@mock.patch(f"{MODULE_PATH}.SubTheme")
128+
@mock.patch(f"{MODULE_PATH}.Theme")
129+
@mock.patch(f"{MODULE_PATH}.transaction.atomic")
130+
@mock.patch(f"{MODULE_PATH}.GeographyType.objects.create")
131+
@mock.patch(f"{MODULE_PATH}.Stratum.objects.create")
132+
@mock.patch(f"{MODULE_PATH}.Age.objects.create")
133+
def test_seed_metrics_data_builds_expected_counts_and_calls(
134+
self,
135+
spy_age_create: mock.MagicMock,
136+
spy_stratum_create: mock.MagicMock,
137+
spy_geography_type_create: mock.MagicMock,
138+
spy_atomic: mock.MagicMock,
139+
spy_theme: mock.MagicMock,
140+
spy_sub_theme: mock.MagicMock,
141+
spy_topic: mock.MagicMock,
142+
spy_metric: mock.MagicMock,
143+
spy_geography: mock.MagicMock,
144+
spy_bulk_create: mock.MagicMock,
145+
spy_seed_time_series_rows: mock.MagicMock,
146+
spy_truncate: mock.MagicMock,
147+
):
148+
spy_atomic.return_value = nullcontext()
149+
spy_theme.side_effect = lambda **kwargs: SimpleNamespace(**kwargs)
150+
spy_sub_theme.side_effect = lambda **kwargs: SimpleNamespace(**kwargs)
151+
spy_topic.side_effect = lambda **kwargs: SimpleNamespace(**kwargs)
152+
spy_metric.side_effect = lambda **kwargs: SimpleNamespace(**kwargs)
153+
spy_geography.side_effect = lambda **kwargs: SimpleNamespace(**kwargs)
154+
spy_geography_type_create.return_value = SimpleNamespace(name="Nation")
155+
spy_stratum_create.return_value = SimpleNamespace(name="All")
156+
spy_age_create.return_value = SimpleNamespace(name="All ages")
157+
spy_seed_time_series_rows.return_value = (77, 88)
158+
159+
themes = [SimpleNamespace(name=f"Theme {index + 1}") for index in range(3)]
160+
sub_themes = [
161+
SimpleNamespace(
162+
name=f"SubTheme {index + 1}", theme=themes[index % len(themes)]
163+
)
164+
for index in range(6)
165+
]
166+
topics = [
167+
SimpleNamespace(
168+
name=f"Topic {index + 1}",
169+
sub_theme=sub_themes[index % len(sub_themes)],
170+
)
171+
for index in range(12)
172+
]
173+
metrics = [
174+
SimpleNamespace(
175+
name=f"Metric {index + 1}", topic=topics[index % len(topics)]
176+
)
177+
for index in range(4)
178+
]
179+
geographies = [
180+
SimpleNamespace(
181+
name=f"Area {index + 1}",
182+
geography_code=f"RND{index + 1:04d}",
183+
geography_type=spy_geography_type_create.return_value,
184+
)
185+
for index in range(2)
186+
]
187+
spy_bulk_create.side_effect = [themes, sub_themes, topics, metrics, geographies]
188+
189+
result = Command._seed_metrics_data(
190+
scale_config={"geographies": 2, "metrics": 4, "days": 9},
191+
truncate_first=True,
192+
)
193+
194+
assert result == {
195+
"Theme": 3,
196+
"SubTheme": 6,
197+
"Topic": 12,
198+
"Metric": 4,
199+
"Geography": 2,
200+
"CoreTimeSeries": 77,
201+
"APITimeSeries": 88,
202+
}
203+
spy_truncate.assert_called_once_with()
204+
spy_seed_time_series_rows.assert_called_once_with(
205+
metrics=metrics,
206+
geographies=geographies,
207+
stratum=spy_stratum_create.return_value,
208+
age=spy_age_create.return_value,
209+
days=9,
210+
)
211+
212+
def test_truncate_metrics_data_deletes_from_all_models(self):
213+
model_names = [
214+
"APITimeSeries",
215+
"CoreTimeSeries",
216+
"Metric",
217+
"Topic",
218+
"SubTheme",
219+
"Theme",
220+
"Geography",
221+
"GeographyType",
222+
"Age",
223+
"Stratum",
224+
]
225+
226+
managers: dict[str, mock.MagicMock] = {}
227+
with ExitStack() as stack:
228+
for model_name in model_names:
229+
manager = mock.MagicMock()
230+
managers[model_name] = manager
231+
stack.enter_context(
232+
mock.patch(f"{MODULE_PATH}.{model_name}.objects", manager)
233+
)
234+
235+
Command._truncate_metrics_data()
236+
237+
for model_name in model_names:
238+
managers[model_name].all.assert_called_once_with()
239+
managers[model_name].all.return_value.delete.assert_called_once_with()
240+
241+
@mock.patch(f"{MODULE_PATH}.APITimeSeries")
242+
@mock.patch(f"{MODULE_PATH}.CoreTimeSeries")
243+
def test_seed_time_series_rows_flushes_remainder(
244+
self,
245+
spy_core_time_series: mock.MagicMock,
246+
spy_api_time_series: mock.MagicMock,
247+
):
248+
spy_core_time_series.side_effect = lambda **kwargs: kwargs
249+
spy_api_time_series.side_effect = lambda **kwargs: kwargs
250+
251+
core_count, api_count = Command._seed_time_series_rows(
252+
metrics=[_fake_metric_hierarchy()],
253+
geographies=[_fake_geography()],
254+
stratum=SimpleNamespace(name="All"),
255+
age=SimpleNamespace(name="All ages"),
256+
days=1,
257+
)
258+
259+
assert core_count == 1
260+
assert api_count == 1
261+
spy_core_time_series.objects.bulk_create.assert_called_once()
262+
spy_api_time_series.objects.bulk_create.assert_called_once()
263+
264+
@mock.patch(f"{MODULE_PATH}.APITimeSeries")
265+
@mock.patch(f"{MODULE_PATH}.CoreTimeSeries")
266+
def test_seed_time_series_rows_flushes_at_batch_size(
267+
self,
268+
spy_core_time_series: mock.MagicMock,
269+
spy_api_time_series: mock.MagicMock,
270+
):
271+
spy_core_time_series.side_effect = lambda **kwargs: kwargs
272+
spy_api_time_series.side_effect = lambda **kwargs: kwargs
273+
274+
core_count, api_count = Command._seed_time_series_rows(
275+
metrics=[_fake_metric_hierarchy()],
276+
geographies=[_fake_geography()],
277+
stratum=SimpleNamespace(name="All"),
278+
age=SimpleNamespace(name="All ages"),
279+
days=5000,
280+
)
281+
282+
assert core_count == 5000
283+
assert api_count == 5000
284+
spy_core_time_series.objects.bulk_create.assert_called_once()
285+
spy_api_time_series.objects.bulk_create.assert_called_once()
286+
287+
def test_bulk_create_materialises_iterable_and_delegates(self):
288+
class FakeModel:
289+
objects = mock.MagicMock()
290+
291+
def records_generator() -> Iterator[int]:
292+
yield 1
293+
yield 2
294+
295+
FakeModel.objects.bulk_create.return_value = ["created-records"]
296+
297+
result = Command._bulk_create(FakeModel, records_generator())
298+
299+
assert result == ["created-records"]
300+
FakeModel.objects.bulk_create.assert_called_once_with([1, 2])
301+
302+
def test_print_summary_writes_expected_output(self):
303+
command = Command()
304+
command.stdout = mock.MagicMock()
305+
306+
command._print_summary(
307+
dataset="both",
308+
scale="small",
309+
seed=123,
310+
counts={
311+
"Theme": 3,
312+
"SubTheme": 6,
313+
"Topic": 12,
314+
"Metric": 10,
315+
"Geography": 5,
316+
"CoreTimeSeries": 1500,
317+
"APITimeSeries": 1500,
318+
},
319+
runtime_seconds=3.456,
320+
)
321+
322+
expected_lines = [
323+
"",
324+
"Seed random summary:",
325+
" dataset: both",
326+
" scale: small",
327+
" seed used: 123",
328+
" Theme: 3",
329+
" SubTheme: 6",
330+
" Topic: 12",
331+
" Metric: 10",
332+
" Geography: 5",
333+
" CoreTimeSeries: 1500",
334+
" APITimeSeries: 1500",
335+
" runtime seconds: 3.46",
336+
]
337+
actual_lines = [call.args[0] for call in command.stdout.write.call_args_list]
338+
339+
assert actual_lines == expected_lines
340+
341+
342+
def test_add_arguments_rejects_invalid_dataset_value():
343+
parser = CommandParser(prog="manage.py seed_random")
344+
Command().add_arguments(parser)
345+
346+
with pytest.raises(SystemExit):
347+
parser.parse_args(["--dataset", "invalid"])

0 commit comments

Comments
 (0)