Merge pull request #28 from m4rc1e/contextmanager

Add context manager to Local

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()

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)