|
5 | 5 | """ |
6 | 6 | Integration tests for delayed saving. |
7 | 7 | """ |
| 8 | +import re |
8 | 9 | import warnings |
9 | 10 |
|
10 | 11 | from cf_units import Unit |
@@ -197,19 +198,36 @@ def test_scheduler_types( |
197 | 198 |
|
198 | 199 | if not save_is_delayed: |
199 | 200 | assert result is None |
200 | | - assert len(logged_warnings) == 2 |
201 | 201 | issued_warnings = [log.message for log in logged_warnings] |
202 | 202 | else: |
203 | 203 | assert result is not None |
204 | 204 | assert len(logged_warnings) == 0 |
205 | | - warnings.simplefilter("error") |
206 | | - issued_warnings = result.compute() |
| 205 | + with warnings.catch_warnings(record=True) as logged_warnings: |
| 206 | + # The compute *returns* warnings from the delayed operations. |
| 207 | + issued_warnings = result.compute() |
| 208 | + issued_warnings = [ |
| 209 | + log.message for log in logged_warnings |
| 210 | + ] + issued_warnings |
| 211 | + |
| 212 | + warning_messages = [warning.args[0] for warning in issued_warnings] |
| 213 | + if scheduler_type == "DistributedScheduler": |
| 214 | + # Ignore any "large data transfer" messages generated, |
| 215 | + # specifically when testing with the Distributed scheduler. |
| 216 | + # These may not always occur and don't reflect something we want to |
| 217 | + # test for. |
| 218 | + large_transfer_message_regex = re.compile( |
| 219 | + "Sending large graph.* may cause some slowdown", re.DOTALL |
| 220 | + ) |
| 221 | + warning_messages = [ |
| 222 | + message |
| 223 | + for message in warning_messages |
| 224 | + if not large_transfer_message_regex.search(message) |
| 225 | + ] |
207 | 226 |
|
208 | | - assert len(issued_warnings) == 2 |
| 227 | + # In all cases, should get 2 fill value warnings overall. |
| 228 | + assert len(warning_messages) == 2 |
209 | 229 | expected_msg = "contains unmasked data points equal to the fill-value" |
210 | | - assert all( |
211 | | - expected_msg in warning.args[0] for warning in issued_warnings |
212 | | - ) |
| 230 | + assert all(expected_msg in message for message in warning_messages) |
213 | 231 |
|
214 | 232 | def test_time_of_writing( |
215 | 233 | self, save_is_delayed, output_path, scheduler_type |
|
0 commit comments