|
12 | 12 | # See the License for the specific language governing permissions and
|
13 | 13 | # limitations under the License.
|
14 | 14 |
|
| 15 | +import re |
15 | 16 | import tempfile
|
16 | 17 |
|
17 | 18 | import pandas as pd
|
18 | 19 |
|
19 | 20 | import bigframes as bf
|
20 | 21 | import bigframes.formatting_helpers as formatting_helpers
|
21 | 22 |
|
| 23 | +job_load_message_regex = r"\w+ job [\w-]+ is \w+\." |
| 24 | + |
22 | 25 |
|
23 | 26 | def test_progress_bar_dataframe(
|
24 | 27 | penguins_df_default_index: bf.dataframe.DataFrame, capsys
|
25 | 28 | ):
|
26 |
| - bf.options.display.progress_bar = "notebook" |
| 29 | + bf.options.display.progress_bar = "terminal" |
| 30 | + capsys.readouterr() # clear output |
27 | 31 | penguins_df_default_index.to_pandas()
|
28 |
| - html_check = "HTML(value=" |
29 |
| - open_job_check = "Open Job" |
30 |
| - lines = capsys.readouterr().out.split("\n") |
31 |
| - lines = [line for line in lines if len(line) > 0] |
32 |
| - assert len(lines) > 0 |
| 32 | + |
| 33 | + assert_loading_msg_exist(capsys.readouterr().out) |
33 | 34 | assert penguins_df_default_index.query_job is not None
|
34 |
| - for line in lines: |
35 |
| - assert html_check in line and open_job_check in line |
36 | 35 |
|
37 | 36 |
|
38 | 37 | def test_progress_bar_series(penguins_df_default_index: bf.dataframe.DataFrame, capsys):
|
39 |
| - bf.options.display.progress_bar = "notebook" |
| 38 | + bf.options.display.progress_bar = "terminal" |
40 | 39 | series = penguins_df_default_index["body_mass_g"].head(10)
|
| 40 | + capsys.readouterr() # clear output |
41 | 41 | series.to_pandas()
|
42 |
| - html_check = "HTML(value=" |
43 |
| - open_job_check = "Open Job" |
44 |
| - lines = capsys.readouterr().out.split("\n") |
45 |
| - lines = [line for line in lines if len(line) > 0] |
46 |
| - assert len(lines) > 0 |
| 42 | + |
| 43 | + assert_loading_msg_exist(capsys.readouterr().out) |
47 | 44 | assert series.query_job is not None
|
48 |
| - for line in lines: |
49 |
| - assert html_check in line and open_job_check in line |
50 | 45 |
|
51 | 46 |
|
52 | 47 | def test_progress_bar_scalar(penguins_df_default_index: bf.dataframe.DataFrame, capsys):
|
53 |
| - bf.options.display.progress_bar = "notebook" |
| 48 | + bf.options.display.progress_bar = "terminal" |
| 49 | + capsys.readouterr() # clear output |
54 | 50 | penguins_df_default_index["body_mass_g"].head(10).mean()
|
55 |
| - html_check = "HTML(value=" |
56 |
| - open_job_check = "Open Job" |
57 |
| - lines = capsys.readouterr().out.split("\n") |
58 |
| - lines = [line for line in lines if len(line) > 0] |
59 |
| - assert len(lines) > 0 |
60 |
| - for line in lines: |
61 |
| - assert html_check in line and open_job_check in line |
| 51 | + |
| 52 | + assert_loading_msg_exist(capsys.readouterr().out) |
62 | 53 |
|
63 | 54 |
|
64 | 55 | def test_progress_bar_read_gbq(session: bf.Session, penguins_table_id: str, capsys):
|
65 |
| - bf.options.display.progress_bar = "notebook" |
| 56 | + bf.options.display.progress_bar = "terminal" |
| 57 | + capsys.readouterr() # clear output |
66 | 58 | session.read_gbq(penguins_table_id)
|
67 |
| - html_check = "HTML(value=" |
68 |
| - open_job_check = "Open Job" |
69 |
| - lines = capsys.readouterr().out.split("\n") |
70 |
| - lines = [line for line in lines if len(line) > 0] |
71 |
| - assert len(lines) > 0 |
72 |
| - for line in lines: |
73 |
| - assert html_check in line and open_job_check in line |
| 59 | + |
| 60 | + assert_loading_msg_exist(capsys.readouterr().out) |
74 | 61 |
|
75 | 62 |
|
76 | 63 | def test_progress_bar_extract_jobs(
|
77 | 64 | penguins_df_default_index: bf.dataframe.DataFrame, gcs_folder, capsys
|
78 | 65 | ):
|
79 |
| - bf.options.display.progress_bar = "notebook" |
| 66 | + bf.options.display.progress_bar = "terminal" |
80 | 67 | path = gcs_folder + "test_read_csv_progress_bar*.csv"
|
| 68 | + capsys.readouterr() # clear output |
81 | 69 | penguins_df_default_index.to_csv(path)
|
82 |
| - html_check = "HTML(value=" |
83 |
| - open_job_check = "Open Job" |
84 |
| - lines = capsys.readouterr().out.split("\n") |
85 |
| - lines = [line for line in lines if len(line) > 0] |
86 |
| - assert len(lines) > 0 |
87 |
| - for line in lines: |
88 |
| - assert html_check in line and open_job_check in line |
| 70 | + |
| 71 | + assert_loading_msg_exist(capsys.readouterr().out) |
89 | 72 |
|
90 | 73 |
|
91 | 74 | def test_progress_bar_load_jobs(
|
92 | 75 | session: bf.Session, penguins_pandas_df_default_index: pd.DataFrame, capsys
|
93 | 76 | ):
|
94 |
| - bf.options.display.progress_bar = "notebook" |
| 77 | + bf.options.display.progress_bar = "terminal" |
95 | 78 | with tempfile.TemporaryDirectory() as dir:
|
96 | 79 | path = dir + "/test_read_csv_progress_bar*.csv"
|
97 | 80 | penguins_pandas_df_default_index.to_csv(path, index=False)
|
| 81 | + capsys.readouterr() # clear output |
98 | 82 | session.read_csv(path)
|
99 |
| - html_check = "HTML(value=" |
100 |
| - open_job_check = "Open Job" |
101 |
| - lines = capsys.readouterr().out.split("\n") |
| 83 | + |
| 84 | + assert_loading_msg_exist(capsys.readouterr().out) |
| 85 | + |
| 86 | + |
| 87 | +def assert_loading_msg_exist(capystOut: str, pattern=job_load_message_regex): |
| 88 | + numLoadingMsg = 0 |
| 89 | + lines = capystOut.split("\n") |
102 | 90 | lines = [line for line in lines if len(line) > 0]
|
| 91 | + |
103 | 92 | assert len(lines) > 0
|
104 | 93 | for line in lines:
|
105 |
| - assert html_check in line and open_job_check in line |
| 94 | + if re.match(pattern, line) is not None: |
| 95 | + numLoadingMsg += 1 |
| 96 | + assert numLoadingMsg > 0 |
106 | 97 |
|
107 | 98 |
|
108 | 99 | def test_query_job_repr_html(penguins_df_default_index: bf.dataframe.DataFrame):
|
109 |
| - bf.options.display.progress_bar = "notebook" |
| 100 | + bf.options.display.progress_bar = "terminal" |
110 | 101 | penguins_df_default_index._block._expr._session.bqclient.default_query_job_config.use_query_cache = (
|
111 | 102 | False
|
112 | 103 | )
|
|
0 commit comments