Fix display handling

1 files changed, 4 insertions, 7 deletions

diff --git a/devbot/xvfb.py b/devbot/xvfb.py
index 6549b9d..eed1468 100644
--- a/devbot/xvfb.py
+++ b/devbot/xvfb.py
@@ -1,11 +1,12 @@
 import os
 import subprocess
 
-from devbot import utils
+from gi.repository import SugarRunner
 
-xvfb_display = ":100"
+from devbot import utils
 
 def start():
+    xvfb_display = SugarRunner.find_free_display()
     xvfb_proc = subprocess.Popen(args=["Xvfb", xvfb_display],
                                  stdout=utils.devnull,
                                  stderr=subprocess.STDOUT)
@@ -15,9 +16,5 @@ def start():
     return (xvfb_proc, orig_display)
 
 def stop(xvfb_proc, orig_display):
-    if orig_display:
-        os.environ["DISPLAY"] = xvfb_display
-    else:
-        del os.environ["DISPLAY"]
-
+    os.environ["DISPLAY"] = orig_display
     xvfb_proc.terminate()