Skip to content

Commit 3d4fc32

Browse files
treysptreysp
committed
Fix jupyter test
1 parent 7fa0eab commit 3d4fc32

1 file changed

Lines changed: 206 additions & 6 deletions

File tree

‎tests/integrations/jupyter/test_magics.py‎

Lines changed: 206 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,9 @@
2222
SUSHI_EXAMPLE_PATH = pathlib.Path("./examples/sushi")
2323
SUCCESS_STYLE = "color: #008000; text-decoration-color: #008000"
2424
NEUTRAL_STYLE = "color: #008080; text-decoration-color: #008080"
25+
BOLD_ONLY = "font-weight: bold"
26+
BOLD_NEUTRAL_STYLE = f"{NEUTRAL_STYLE}; {BOLD_ONLY}"
27+
BOLD_SUCCESS_STYLE = f"{SUCCESS_STYLE}; {BOLD_ONLY}"
2528
RICH_PRE_STYLE = "white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"
2629
FREEZE_TIME = "2023-01-01 00:00:00 UTC"
2730

@@ -289,9 +292,6 @@ def test_plan(
289292
with capture_output() as output:
290293
notebook.run_line_magic(magic_name="plan", line="--no-prompts --auto-apply")
291294

292-
# TODO: Should this be going to stdout? This is printing the status updates for when each batch finishes for
293-
# the models and how long it took
294-
assert len(output.stdout.strip().split("\n")) == 46
295295
assert not output.stderr
296296
assert len(output.outputs) == 4
297297
text_output = convert_all_html_output_to_text(output)
@@ -321,15 +321,215 @@ def test_run_dag(
321321
notebook.run_line_magic(magic_name="run_dag", line="")
322322

323323
assert not output.stdout.startswith(
324-
"'Evaluating models ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% • 18/18"
324+
"'Executing model batches ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% • 18/18"
325325
)
326326
assert not output.stderr
327327
assert len(output.outputs) == 6
328-
assert convert_all_html_output_to_text(output) == [
328+
html_text_actual = convert_all_html_output_to_text(output)
329+
html_text_expected = [
330+
"[2K",
331+
"Executing model batches ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸━━ 93.8% • 15/16 • 0:00:",
332+
"[1/1] sushi.waiter_as_customer_by_day [insert 2023-01-01 - 2023-01-02, audits ✔2]",
333+
"",
329334
"✔ Model batches executed",
330335
"Run finished for environment 'prod'",
331336
]
332-
assert get_all_html_output(output) == [
337+
for txt_actual, txt_expected in zip(html_text_actual, html_text_expected):
338+
assert txt_actual.startswith(txt_expected)
339+
actual_html_output = get_all_html_output(output)
340+
# Replace dynamic elapsed time with 00
341+
for i, chunk in enumerate(actual_html_output):
342+
pattern = 'font\-weight: bold">0\.</span>\\d{2}s </pre>'
343+
import re
344+
345+
actual_html_output[i] = re.sub(pattern, 'font-weight: bold">0.</span>00s </pre>', chunk)
346+
expected_html_output = [
347+
str(
348+
h(
349+
"pre",
350+
{"style": RICH_PRE_STYLE},
351+
"\x1b",
352+
h(
353+
"span",
354+
{"style": BOLD_ONLY},
355+
"[",
356+
autoescape=False,
357+
),
358+
"2K",
359+
autoescape=False,
360+
)
361+
),
362+
str(
363+
h(
364+
"pre",
365+
{"style": RICH_PRE_STYLE},
366+
h(
367+
"span",
368+
{"style": "color: #000080; text-decoration-color: #000080; font-weight: bold"},
369+
"Executing model batches",
370+
autoescape=False,
371+
),
372+
" ",
373+
h(
374+
"span",
375+
{"style": "color: #f92672; text-decoration-color: #f92672"},
376+
"━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸",
377+
autoescape=False,
378+
),
379+
h(
380+
"span",
381+
{"style": "color: #3a3a3a; text-decoration-color: #3a3a3a"},
382+
"━━",
383+
autoescape=False,
384+
),
385+
" ",
386+
h(
387+
"span",
388+
{"style": "color: #800080; text-decoration-color: #800080"},
389+
"93.8%",
390+
autoescape=False,
391+
),
392+
" • ",
393+
h(
394+
"span",
395+
{"style": SUCCESS_STYLE},
396+
"15/16",
397+
autoescape=False,
398+
),
399+
" • ",
400+
h(
401+
"span",
402+
{"style": "color: #808000; text-decoration-color: #808000"},
403+
"0:00:00",
404+
autoescape=False,
405+
),
406+
"sushi.waiter_as_customer_by_day ",
407+
h(
408+
"span",
409+
{"style": SUCCESS_STYLE},
410+
".. ",
411+
autoescape=False,
412+
),
413+
" ",
414+
autoescape=False,
415+
)
416+
),
417+
str(
418+
h(
419+
"pre",
420+
{"style": RICH_PRE_STYLE},
421+
h(
422+
"span",
423+
{"style": BOLD_ONLY},
424+
"[",
425+
autoescape=False,
426+
),
427+
h(
428+
"span",
429+
{"style": BOLD_NEUTRAL_STYLE},
430+
"1",
431+
autoescape=False,
432+
),
433+
"/",
434+
h(
435+
"span",
436+
{"style": BOLD_NEUTRAL_STYLE},
437+
"1",
438+
autoescape=False,
439+
),
440+
h(
441+
"span",
442+
{"style": BOLD_ONLY},
443+
"]",
444+
autoescape=False,
445+
),
446+
" sushi.waiter_as_customer_by_day ",
447+
h(
448+
"span",
449+
{"style": BOLD_ONLY},
450+
"[",
451+
autoescape=False,
452+
),
453+
"insert ",
454+
h(
455+
"span",
456+
{"style": BOLD_NEUTRAL_STYLE},
457+
"2023",
458+
autoescape=False,
459+
),
460+
"-",
461+
h(
462+
"span",
463+
{"style": BOLD_NEUTRAL_STYLE},
464+
"01",
465+
autoescape=False,
466+
),
467+
"-",
468+
h(
469+
"span",
470+
{"style": BOLD_NEUTRAL_STYLE},
471+
"01",
472+
autoescape=False,
473+
),
474+
" - ",
475+
h(
476+
"span",
477+
{"style": BOLD_NEUTRAL_STYLE},
478+
"2023",
479+
autoescape=False,
480+
),
481+
"-",
482+
h(
483+
"span",
484+
{"style": BOLD_NEUTRAL_STYLE},
485+
"01",
486+
autoescape=False,
487+
),
488+
"-",
489+
h(
490+
"span",
491+
{"style": BOLD_NEUTRAL_STYLE},
492+
"02",
493+
autoescape=False,
494+
),
495+
", audits ",
496+
h(
497+
"span",
498+
{"style": SUCCESS_STYLE},
499+
"✔",
500+
autoescape=False,
501+
),
502+
h(
503+
"span",
504+
{"style": BOLD_NEUTRAL_STYLE},
505+
"2",
506+
autoescape=False,
507+
),
508+
h(
509+
"span",
510+
{"style": BOLD_ONLY},
511+
"]",
512+
autoescape=False,
513+
),
514+
" ",
515+
h(
516+
"span",
517+
{"style": BOLD_NEUTRAL_STYLE},
518+
"0.",
519+
autoescape=False,
520+
),
521+
"00s ",
522+
autoescape=False,
523+
)
524+
),
525+
str(
526+
h(
527+
"pre",
528+
{"style": RICH_PRE_STYLE},
529+
"",
530+
autoescape=False,
531+
)
532+
),
333533
str(
334534
h(
335535
"pre",

0 commit comments

Comments
 (0)