When os.environ has no DISPLAY variable, PyVirtualDisplay sets it to ":0" when calling stop() instead of deleting it

[adveres]

Hi,

First, thank you for the great library. I believe I've found a bug. You fixed Issue #2 by adding a default DISPLAY of ":0" when stopping the display. This is an issue if we try to create multiple displays sequentially, because stop() is incorrectly leaving the DISPLAY environment variable set, when the key should be deleted. I have the same use-case as Issue #2, where the CentOS machine I am running on does not have a DISPLAY environment variable by default.

See script/output below, and let me know if you need any more details. I am using xvfb on a CentOS server that has no DISPLAY by default.

import os
import time
from pyvirtualdisplay import Display
from selenium import webdriver
from selenium.webdriver.firefox.firefox_binary import FirefoxBinary

def start():
    display = None

    if 'DISPLAY' not in os.environ:
        print("  Bringing up a new display, because os.environ does not have one.")
        display = Display(visible=False, size=(1920, 1200))  # TODO is this cheating?? Huge resolution.
        display.start()

        print("  Started pyvirtualdisplay with backend '%s', screen '%s'." % (display.backend, display.screen))
    else:
        print("  Environ DISPLAY is '%s', no need to use pyvirtualdisplay." % os.environ['DISPLAY'])

    # TODO do we need to sleep here for the display to start?

    try:
        with open("/tmp/firefoxLog.txt", 'w') as log_file:
            binary = FirefoxBinary(firefox_path='/usr/bin/firefox', log_file=log_file)
            browser = webdriver.Firefox(firefox_binary=binary)
            browser.get("http://www.google.com")
    except:
        if display:
            print("  Stopping display because the browser failed to come up.")
            display.stop()
        else:
            print("  No need to stop display, because we did not make one earlier.")

        raise

    return display, browser

def stop(display, browser):
    if browser:
        print "  Quitting browser"
        browser.quit()
    if display:
        print "  Stopping display"
        display.stop()

print "TEST1"
d1, b1 = start()
time.sleep(3)
stop(d1, b1)

time.sleep(2)

print "\nTEST2"
d2, b2 = start()
time.sleep(3)
stop(d2, b2)

exit(0)

I believe this can be fixed by deleting "DISPLAY" from the environ dictionary if it was not there when we initially started.

Thank you!

[ponty]

Thanks for the report!