Skip to content

Commit aa9a5b5

Browse files
izeigermanizeigerman
authored
Chore: Add more integration tests (#1563)
1 parent b7275d0 commit aa9a5b5

1 file changed

Lines changed: 367 additions & 0 deletions

File tree

‎tests/core/test_integration.py‎

Lines changed: 367 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import numpy as np
88
import pandas as pd
99
import pytest
10+
from freezegun import freeze_time
1011
from pytest_mock.plugin import MockerFixture
1112
from sqlglot import exp
1213
from sqlglot.expressions import DataType
@@ -56,6 +57,364 @@ def plan_choice(plan: Plan, choice: SnapshotChangeCategory) -> None:
5657
plan.set_choice(snapshot, choice)
5758

5859

60+
@freeze_time("2023-01-08 15:00:00")
61+
@pytest.mark.integration
62+
@pytest.mark.core_integration
63+
def test_forward_only_plan_with_effective_date(mocker: MockerFixture):
64+
context, plan = init_and_plan_context("examples/sushi", mocker)
65+
context.apply(plan)
66+
67+
model_name = "sushi.waiter_revenue_by_day"
68+
model = context.models[model_name]
69+
context.upsert_model(add_projection_to_model(t.cast(SqlModel, model)))
70+
71+
plan = context.plan("dev", no_prompts=True, skip_tests=True, forward_only=True)
72+
assert len(plan.new_snapshots) == 2
73+
assert (
74+
plan.context_diff.snapshots[model_name].change_category
75+
== SnapshotChangeCategory.FORWARD_ONLY
76+
)
77+
assert (
78+
plan.context_diff.snapshots["sushi.top_waiters"].change_category
79+
== SnapshotChangeCategory.FORWARD_ONLY
80+
)
81+
assert plan.start == to_date("2023-01-07")
82+
assert plan.missing_intervals == [
83+
SnapshotIntervals(
84+
snapshot_name=model_name,
85+
intervals=[(to_timestamp("2023-01-07"), to_timestamp("2023-01-08"))],
86+
),
87+
SnapshotIntervals(
88+
snapshot_name="sushi.top_waiters",
89+
intervals=[(to_timestamp("2023-01-07"), to_timestamp("2023-01-08"))],
90+
),
91+
]
92+
93+
plan.effective_from = "2023-01-05"
94+
# Default start should be set to effective_from
95+
assert plan.missing_intervals == [
96+
SnapshotIntervals(
97+
snapshot_name=model_name,
98+
intervals=[
99+
(to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
100+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
101+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
102+
],
103+
),
104+
SnapshotIntervals(
105+
snapshot_name="sushi.top_waiters",
106+
intervals=[
107+
(to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
108+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
109+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
110+
],
111+
),
112+
]
113+
114+
plan.start = "2023-01-06"
115+
# Start override should take precedence
116+
assert plan.missing_intervals == [
117+
SnapshotIntervals(
118+
snapshot_name=model_name,
119+
intervals=[
120+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
121+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
122+
],
123+
),
124+
SnapshotIntervals(
125+
snapshot_name="sushi.top_waiters",
126+
intervals=[
127+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
128+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
129+
],
130+
),
131+
]
132+
133+
plan.effective_from = "2023-01-04"
134+
# Start should remain unchanged
135+
assert plan.start == "2023-01-06"
136+
assert plan.missing_intervals == [
137+
SnapshotIntervals(
138+
snapshot_name=model_name,
139+
intervals=[
140+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
141+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
142+
],
143+
),
144+
SnapshotIntervals(
145+
snapshot_name="sushi.top_waiters",
146+
intervals=[
147+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
148+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
149+
],
150+
),
151+
]
152+
153+
context.apply(plan)
154+
155+
dev_df = context.engine_adapter.fetchdf(
156+
"SELECT DISTINCT ds FROM sushi__dev.waiter_revenue_by_day ORDER BY ds"
157+
)
158+
assert dev_df["ds"].tolist() == ["2023-01-06", "2023-01-07"]
159+
160+
prod_plan = context.plan(no_prompts=True, skip_tests=True)
161+
# Make sure that the previously set effective_from is respected
162+
assert prod_plan.start == to_timestamp("2023-01-04")
163+
assert prod_plan.missing_intervals == [
164+
SnapshotIntervals(
165+
snapshot_name=model_name,
166+
intervals=[
167+
(to_timestamp("2023-01-04"), to_timestamp("2023-01-05")),
168+
(to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
169+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
170+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
171+
],
172+
),
173+
SnapshotIntervals(
174+
snapshot_name="sushi.top_waiters",
175+
intervals=[
176+
(to_timestamp("2023-01-04"), to_timestamp("2023-01-05")),
177+
(to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
178+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
179+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
180+
],
181+
),
182+
]
183+
184+
context.apply(prod_plan)
185+
186+
prod_df = context.engine_adapter.fetchdf(
187+
"SELECT DISTINCT ds FROM sushi.waiter_revenue_by_day WHERE one IS NOT NULL ORDER BY ds"
188+
)
189+
assert prod_df["ds"].tolist() == ["2023-01-04", "2023-01-05", "2023-01-06", "2023-01-07"]
190+
191+
192+
@freeze_time("2023-01-08 15:00:00")
193+
@pytest.mark.integration
194+
@pytest.mark.core_integration
195+
def test_forward_only_model_regular_plan(mocker: MockerFixture):
196+
context, plan = init_and_plan_context("examples/sushi", mocker)
197+
context.apply(plan)
198+
199+
model_name = "sushi.waiter_revenue_by_day"
200+
201+
model = context.models[model_name]
202+
model = add_projection_to_model(t.cast(SqlModel, model))
203+
forward_only_kind = model.kind.copy(update={"forward_only": True})
204+
model = model.copy(update={"kind": forward_only_kind})
205+
206+
context.upsert_model(model)
207+
208+
plan = context.plan("dev", no_prompts=True, skip_tests=True)
209+
assert len(plan.new_snapshots) == 2
210+
assert (
211+
plan.context_diff.snapshots[model_name].change_category
212+
== SnapshotChangeCategory.FORWARD_ONLY
213+
)
214+
assert (
215+
plan.context_diff.snapshots["sushi.top_waiters"].change_category
216+
== SnapshotChangeCategory.FORWARD_ONLY
217+
)
218+
assert plan.start == to_datetime("2023-01-01")
219+
assert not plan.missing_intervals
220+
221+
context.apply(plan)
222+
223+
dev_df = context.engine_adapter.fetchdf(
224+
"SELECT DISTINCT ds FROM sushi__dev.waiter_revenue_by_day ORDER BY ds"
225+
)
226+
assert not dev_df["ds"].tolist()
227+
228+
# Run a restatement plan to preview changes
229+
plan = context.plan("dev", no_prompts=True, skip_tests=True, restate_models=[model_name])
230+
plan.start = "2023-01-06"
231+
assert plan.missing_intervals == [
232+
SnapshotIntervals(
233+
snapshot_name=model_name,
234+
intervals=[
235+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
236+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
237+
],
238+
),
239+
SnapshotIntervals(
240+
snapshot_name="sushi.top_waiters",
241+
intervals=[
242+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
243+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
244+
],
245+
),
246+
]
247+
248+
# Make sure that changed start is reflected in missing intervals
249+
plan.start = "2023-01-07"
250+
assert plan.missing_intervals == [
251+
SnapshotIntervals(
252+
snapshot_name=model_name,
253+
intervals=[
254+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
255+
],
256+
),
257+
SnapshotIntervals(
258+
snapshot_name="sushi.top_waiters",
259+
intervals=[
260+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
261+
],
262+
),
263+
]
264+
265+
context.apply(plan)
266+
267+
dev_df = context.engine_adapter.fetchdf(
268+
"SELECT DISTINCT ds FROM sushi__dev.waiter_revenue_by_day ORDER BY ds"
269+
)
270+
assert dev_df["ds"].tolist() == ["2023-01-07"]
271+
272+
# Promote changes to prod
273+
prod_plan = context.plan(no_prompts=True, skip_tests=True)
274+
assert not prod_plan.missing_intervals
275+
276+
context.apply(prod_plan)
277+
278+
# The change was applied in a forward-only manner so no values in the new column should be populated
279+
prod_df = context.engine_adapter.fetchdf(
280+
"SELECT DISTINCT ds FROM sushi.waiter_revenue_by_day WHERE one IS NOT NULL ORDER BY ds"
281+
)
282+
assert not prod_df["ds"].tolist()
283+
284+
285+
@freeze_time("2023-01-08 15:00:00")
286+
@pytest.mark.integration
287+
@pytest.mark.core_integration
288+
def test_plan_set_choice_is_reflected_in_missing_intervals(mocker: MockerFixture):
289+
context, plan = init_and_plan_context("examples/sushi", mocker)
290+
context.apply(plan)
291+
292+
model_name = "sushi.waiter_revenue_by_day"
293+
294+
model = context.models[model_name]
295+
context.upsert_model(add_projection_to_model(t.cast(SqlModel, model)))
296+
297+
plan = context.plan("dev", no_prompts=True, skip_tests=True)
298+
assert len(plan.new_snapshots) == 2
299+
assert (
300+
plan.context_diff.snapshots[model_name].change_category
301+
== SnapshotChangeCategory.NON_BREAKING
302+
)
303+
assert (
304+
plan.context_diff.snapshots["sushi.top_waiters"].change_category
305+
== SnapshotChangeCategory.INDIRECT_NON_BREAKING
306+
)
307+
assert plan.start == to_timestamp("2023-01-01")
308+
assert plan.missing_intervals == [
309+
SnapshotIntervals(
310+
snapshot_name=model_name,
311+
intervals=[
312+
(to_timestamp("2023-01-01"), to_timestamp("2023-01-02")),
313+
(to_timestamp("2023-01-02"), to_timestamp("2023-01-03")),
314+
(to_timestamp("2023-01-03"), to_timestamp("2023-01-04")),
315+
(to_timestamp("2023-01-04"), to_timestamp("2023-01-05")),
316+
(to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
317+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
318+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
319+
],
320+
),
321+
]
322+
323+
# Change the category to BREAKING
324+
plan.set_choice(plan.context_diff.snapshots[model_name], SnapshotChangeCategory.BREAKING)
325+
assert (
326+
plan.context_diff.snapshots[model_name].change_category == SnapshotChangeCategory.BREAKING
327+
)
328+
assert (
329+
plan.context_diff.snapshots["sushi.top_waiters"].change_category
330+
== SnapshotChangeCategory.INDIRECT_BREAKING
331+
)
332+
assert plan.missing_intervals == [
333+
SnapshotIntervals(
334+
snapshot_name=model_name,
335+
intervals=[
336+
(to_timestamp("2023-01-01"), to_timestamp("2023-01-02")),
337+
(to_timestamp("2023-01-02"), to_timestamp("2023-01-03")),
338+
(to_timestamp("2023-01-03"), to_timestamp("2023-01-04")),
339+
(to_timestamp("2023-01-04"), to_timestamp("2023-01-05")),
340+
(to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
341+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
342+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
343+
],
344+
),
345+
SnapshotIntervals(
346+
snapshot_name="sushi.top_waiters",
347+
intervals=[
348+
(to_timestamp("2023-01-01"), to_timestamp("2023-01-02")),
349+
(to_timestamp("2023-01-02"), to_timestamp("2023-01-03")),
350+
(to_timestamp("2023-01-03"), to_timestamp("2023-01-04")),
351+
(to_timestamp("2023-01-04"), to_timestamp("2023-01-05")),
352+
(to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
353+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
354+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
355+
],
356+
),
357+
]
358+
359+
# Change the category back to NON_BREAKING
360+
plan.set_choice(plan.context_diff.snapshots[model_name], SnapshotChangeCategory.NON_BREAKING)
361+
assert (
362+
plan.context_diff.snapshots[model_name].change_category
363+
== SnapshotChangeCategory.NON_BREAKING
364+
)
365+
assert (
366+
plan.context_diff.snapshots["sushi.top_waiters"].change_category
367+
== SnapshotChangeCategory.INDIRECT_NON_BREAKING
368+
)
369+
assert plan.missing_intervals == [
370+
SnapshotIntervals(
371+
snapshot_name=model_name,
372+
intervals=[
373+
(to_timestamp("2023-01-01"), to_timestamp("2023-01-02")),
374+
(to_timestamp("2023-01-02"), to_timestamp("2023-01-03")),
375+
(to_timestamp("2023-01-03"), to_timestamp("2023-01-04")),
376+
(to_timestamp("2023-01-04"), to_timestamp("2023-01-05")),
377+
(to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
378+
(to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
379+
(to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
380+
],
381+
),
382+
]
383+
384+
context.apply(plan)
385+
386+
dev_df = context.engine_adapter.fetchdf(
387+
"SELECT DISTINCT ds FROM sushi__dev.waiter_revenue_by_day ORDER BY ds"
388+
)
389+
assert dev_df["ds"].tolist() == [
390+
"2023-01-01",
391+
"2023-01-02",
392+
"2023-01-03",
393+
"2023-01-04",
394+
"2023-01-05",
395+
"2023-01-06",
396+
"2023-01-07",
397+
]
398+
399+
# Promote changes to prod
400+
prod_plan = context.plan(no_prompts=True, skip_tests=True)
401+
assert not prod_plan.missing_intervals
402+
403+
context.apply(prod_plan)
404+
prod_df = context.engine_adapter.fetchdf(
405+
"SELECT DISTINCT ds FROM sushi.waiter_revenue_by_day WHERE one IS NOT NULL ORDER BY ds"
406+
)
407+
assert prod_df["ds"].tolist() == [
408+
"2023-01-01",
409+
"2023-01-02",
410+
"2023-01-03",
411+
"2023-01-04",
412+
"2023-01-05",
413+
"2023-01-06",
414+
"2023-01-07",
415+
]
416+
417+
59418
@pytest.mark.integration
60419
@pytest.mark.core_integration
61420
@pytest.mark.parametrize(
@@ -1141,3 +1500,11 @@ def start(context: Context) -> TimeLike:
11411500
env = context.state_sync.get_environment("prod")
11421501
assert env
11431502
return env.start_at
1503+
1504+
1505+
def add_projection_to_model(model: SqlModel) -> SqlModel:
1506+
kwargs = {
1507+
**model.dict(),
1508+
"query": model.query.select(exp.Literal.number(1).as_("one")), # type: ignore
1509+
}
1510+
return SqlModel.parse_obj(kwargs)

0 commit comments

Comments
 (0)