From 87b359439f62a79daef71ee3d0eebfca67ffb969 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Wed, 12 Aug 2020 04:54:28 -0700 Subject: [PATCH] [test] replace capsys with the more refined CaptureStderr/CaptureStdout (#6422) * replace capsys with the more refined CaptureStderr/CaptureStdout * Update examples/seq2seq/test_seq2seq_examples.py Co-authored-by: Sam Shleifer --- examples/seq2seq/test_seq2seq_examples.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/examples/seq2seq/test_seq2seq_examples.py b/examples/seq2seq/test_seq2seq_examples.py index 06719446d6..70b34d1d8a 100644 --- a/examples/seq2seq/test_seq2seq_examples.py +++ b/examples/seq2seq/test_seq2seq_examples.py @@ -15,7 +15,7 @@ from torch.utils.data import DataLoader import lightning_base from transformers import AutoModelForSeq2SeqLM, AutoTokenizer -from transformers.testing_utils import require_multigpu +from transformers.testing_utils import CaptureStderr, CaptureStdout, require_multigpu from .distillation import distill_main, evaluate_checkpoint from .finetune import SummarizationModule, main @@ -329,7 +329,7 @@ def test_finetune_extra_model_args(): assert str(excinfo.value) == f"model config doesn't have a `{unsupported_param}` attribute" -def test_finetune_lr_shedulers(capsys): +def test_finetune_lr_schedulers(): args_d: dict = CHEAP_ARGS.copy() task = "summarization" @@ -361,23 +361,23 @@ def test_finetune_lr_shedulers(capsys): # --help test with pytest.raises(SystemExit) as excinfo: - args = parser.parse_args(args) + with CaptureStdout() as cs: + args = parser.parse_args(args) assert False, "--help is expected to sys.exit" assert excinfo.type == SystemExit - captured = capsys.readouterr() expected = lightning_base.arg_to_scheduler_metavar - assert expected in captured.out, "--help is expected to list the supported schedulers" + assert expected in cs.out, "--help is expected to list the supported schedulers" # --lr_scheduler=non_existing_scheduler test unsupported_param = "non_existing_scheduler" args = {f"--lr_scheduler={unsupported_param}"} with pytest.raises(SystemExit) as excinfo: - args = parser.parse_args(args) + with CaptureStderr() as cs: + args = parser.parse_args(args) assert False, "invalid argument is expected to sys.exit" assert excinfo.type == SystemExit - captured = capsys.readouterr() expected = f"invalid choice: '{unsupported_param}'" - assert expected in captured.err, f"should have bailed on invalid choice of scheduler {unsupported_param}" + assert expected in cs.err, f"should have bailed on invalid choice of scheduler {unsupported_param}" # --lr_scheduler=existing_scheduler test supported_param = "cosine"