From 5df90e0991ab471d51af2faa43041d3e60115db5 Mon Sep 17 00:00:00 2001 From: Marc Foley Date: Fri, 4 Dec 2020 11:36:23 +0000 Subject: [PATCH] Add context manager to Local --- browserstack/local.py | 14 +++++++++++--- tests/test_local.py | 4 ++++ 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/browserstack/local.py b/browserstack/local.py index cb3bd2c..43bc237 100644 --- a/browserstack/local.py +++ b/browserstack/local.py @@ -3,9 +3,9 @@ from browserstack.bserrors import BrowserStackLocalError class Local: - def __init__(self, key=None, binary_path=None): + def __init__(self, key=None, binary_path=None, **kwargs): self.key = os.environ['BROWSERSTACK_ACCESS_KEY'] if 'BROWSERSTACK_ACCESS_KEY' in os.environ else key - self.options = None + self.options = kwargs self.local_logfile_path = os.path.join(os.getcwd(), 'local.log') def __xstr(self, key, value): @@ -29,7 +29,8 @@ def _generate_stop_cmd(self): return cmd def start(self, **kwargs): - self.options = kwargs + for k, v in kwargs.items(): + self.options[k] = v if 'key' in self.options: self.key = self.options['key'] @@ -74,3 +75,10 @@ def stop(self): (out, err) = proc.communicate() except Exception as e: return + + def __enter__(self): + self.start(**self.options) + return self + + def __exit__(self, *args): + self.stop() diff --git a/tests/test_local.py b/tests/test_local.py index d6ac57b..3d455d3 100644 --- a/tests/test_local.py +++ b/tests/test_local.py @@ -79,3 +79,7 @@ def test_local_identifier(self): self.local.start(localIdentifier='mytunnel', onlyCommand=True) self.assertIn('-localIdentifier', self.local._generate_cmd()) self.assertIn('mytunnel', self.local._generate_cmd()) + + def test_context_manager(self): + with Local('BROWSERSTACK_ACCESS_KEY') as local: + self.assertNotEqual(local.proc.pid, 0)