From f87d6938556c6676fe0bc69f365f87cf5b3ade62 Mon Sep 17 00:00:00 2001
From: Corey Goldberg <cgoldberg@gmail.com>
Date: Sat, 18 May 2013 11:09:08 -0400
Subject: [PATCH] handle bad args and more unit tests

---
 xvfbwrapper/test_xvfb.py   | 49 ++++++++++++++++++++++++++++++++------
 xvfbwrapper/xvfbwrapper.py | 10 ++++++--
 2 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/xvfbwrapper/test_xvfb.py b/xvfbwrapper/test_xvfb.py
index 762a50f..16a3f85 100644
--- a/xvfbwrapper/test_xvfb.py
+++ b/xvfbwrapper/test_xvfb.py
@@ -4,8 +4,14 @@
 from .xvfbwrapper import Xvfb
 
 import os
+import sys
 import unittest
 
+try:
+    from StringIO import StringIO  # Python 2.7 compat
+except ImportError:
+    from io import StringIO
+
 import pep8
 
 
@@ -18,7 +24,15 @@ class TestXvfb(unittest.TestCase):
         self.assertEqual(':%d' % xvfb.vdisplay_num, os.environ['DISPLAY'])
         self.assertIsNot(None, xvfb.proc)
 
-    def test_start_with_args(self):
+    def test_stop(self):
+        orig_display = os.environ['DISPLAY']
+        xvfb = Xvfb()
+        xvfb.start()
+        self.assertNotEqual(orig_display, os.environ['DISPLAY'])
+        xvfb.stop()
+        self.assertEqual(orig_display, os.environ['DISPLAY'])
+
+    def test_start_with_kwargs(self):
         w = 800
         h = 600
         depth = 16
@@ -29,15 +43,36 @@ class TestXvfb(unittest.TestCase):
         self.assertEqual(h, xvfb.height)
         self.assertEqual(depth, xvfb.colordepth)
         self.assertEqual(os.environ['DISPLAY'], ':%d' % xvfb.vdisplay_num)
+        self.assertIsNotNone(xvfb.proc)
+
+    def test_start_with_arbitrary_kwarg(self):
+        xvfb = Xvfb(nolisten='tcp')
+        self.addCleanup(xvfb.stop)
+        xvfb.start()
+        self.assertEqual(os.environ['DISPLAY'], ':%d' % xvfb.vdisplay_num)
         self.assertIsNot(None, xvfb.proc)
 
-    def test_stop(self):
-        orig = os.environ['DISPLAY']
-        xvfb = Xvfb()
+    def test_start_with_bad_kwarg_param(self):
+        orig_display = os.environ['DISPLAY']
+        out = StringIO()
+        sys.stdout = out
+        xvfb = Xvfb(nolisten='badvalue')
+        self.addCleanup(xvfb.stop)
         xvfb.start()
-        self.assertNotEqual(orig, os.environ['DISPLAY'])
-        xvfb.stop()
-        self.assertEqual(orig, os.environ['DISPLAY'])
+        self.assertEqual(orig_display, os.environ['DISPLAY'])
+        self.assertIsNone(xvfb.proc)
+        self.assertIn('Error: Xvfb did not start', out.getvalue())
+
+    def test_start_with_bad_kwarg(self):
+        orig_display = os.environ['DISPLAY']
+        out = StringIO()
+        sys.stdout = out
+        xvfb = Xvfb(badarg='foo')
+        self.addCleanup(xvfb.stop)
+        xvfb.start()
+        self.assertEqual(orig_display, os.environ['DISPLAY'])
+        self.assertIsNone(xvfb.proc)
+        self.assertIn('Error: Xvfb did not start', out.getvalue())
 
 
 class Pep8ConformanceTestCase(unittest.TestCase):
diff --git a/xvfbwrapper/xvfbwrapper.py b/xvfbwrapper/xvfbwrapper.py
index 01bf4fe..8f3278d 100644
--- a/xvfbwrapper/xvfbwrapper.py
+++ b/xvfbwrapper/xvfbwrapper.py
@@ -43,9 +43,15 @@ class Xvfb:
 
         self.proc = subprocess.Popen(self.xvfb_cmd,
                                      stdout=open(os.devnull),
-                                     stderr=open(os.devnull))
+                                     stderr=open(os.devnull),
+                                     )
         time.sleep(0.1)  # give Xvfb time to start
-        self._redirect_display(self.vdisplay_num)
+        ret_code = self.proc.poll()
+        if ret_code is None:
+            self._redirect_display(self.vdisplay_num)
+        else:
+            self.proc = None
+            print('Error: Xvfb did not start')
 
     def stop(self):
         self._redirect_display(self.old_display_num)