From d62610c44eb7f85ce5ba501b795f7e51715150c7 Mon Sep 17 00:00:00 2001 From: Tyler Goodlet Date: Mon, 10 Feb 2020 12:59:44 -0500 Subject: [PATCH] Search for guard and strip instead of hardcoding --- tests/test_docs_examples.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/test_docs_examples.py b/tests/test_docs_examples.py index f45e6eb..c6d5df8 100644 --- a/tests/test_docs_examples.py +++ b/tests/test_docs_examples.py @@ -2,6 +2,7 @@ Let's make sure them docs work yah? """ from contextlib import contextmanager +import itertools import os import sys import subprocess @@ -42,9 +43,13 @@ def run_example_in_subproc(loglevel, testdir, arb_addr): os.path.join(str(testdir), '__main__.py') ) - # drop the ``if __name__ == '__main__'`` guard from the *NIX - # version of each script - script_code = '\n'.join(script_code.splitlines()[:-4]) + # drop the ``if __name__ == '__main__'`` guard onwards from + # the *NIX version of each script + windows_script_lines = itertools.takewhile( + lambda line: "if __name__ ==" not in line, + script_code.splitlines() + ) + script_code = '\n'.join(windows_script_lines) script_file = testdir.makefile('.py', script_code) # without this, tests hang on windows forever