kea/tests/lettuce/features/terrain/terrain.py

#
# This is the 'terrain' in which the lettuce lives. By convention, this is
# where global setup and teardown is defined.
#
# We declare some attributes of the global 'world' variables here, so the
# tests can safely assume they are present.
#
# We also use it to provide scenario invariants, such as resetting data.
#
from lettuce import *
import subprocess
import os.path
import shutil
import re
import time

# In order to make sure we start all tests with a 'clean' environment,
# We perform a number of initialization steps, like restoring configuration
# files, and removing generated data files.

# This approach may not scale; if so we should probably provide specific
# initialization steps for scenarios. But until that is shown to be a problem,
# It will keep the scenarios cleaner.

# This is a list of files that are freshly copied before each scenario
# The first element is the original, the second is the target that will be
# used by the tests that need them
copylist = [
["configurations/example.org.config.orig", "configurations/example.org.config"]
]

# This is a list of files that, if present, will be removed before a scenario
removelist = [
"data/test_nonexistent_db.sqlite3"
]

# When waiting for output data of a running process, use OUTPUT_WAIT_INTERVAL
# as the interval in which to check again if it has not been found yet.
# If we have waited OUTPUT_WAIT_MAX_INTERVALS times, we will abort with an
# error (so as not to hang indefinitely)
OUTPUT_WAIT_INTERVAL = 0.5
OUTPUT_WAIT_MAX_INTERVALS = 20

# class that keeps track of one running process and the files
# we created for it.
class RunningProcess:
    def __init__(self, step, process_name, args):
        # set it to none first so destructor won't error if initializer did
        self.process = None
        self.step = step
        self.process_name = process_name
        self.remove_files_on_exit = True
        self._check_output_dir()
        self._create_filenames()
        self._start_process(args)

    def _start_process(self, args):
        stderr_write = open(self.stderr_filename, "w")
        stdout_write = open(self.stdout_filename, "w")
        self.process = subprocess.Popen(args, 1, None, subprocess.PIPE,
                                        stdout_write, stderr_write)
        # open them again, this time for reading
        self.stderr = open(self.stderr_filename, "r")
        self.stdout = open(self.stdout_filename, "r")

    def mangle_filename(self, filebase, extension):
        filebase = re.sub("\s+", "_", filebase)
        filebase = re.sub("[^a-zA-Z0-9.\-_]", "", filebase)
        return filebase + "." + extension

    def _check_output_dir(self):
        # We may want to make this overridable by the user, perhaps
        # through an environment variable. Since we currently expect
        # lettuce to be run from our lettuce dir, we shall just use
        # the relative path 'output/'
        self._output_dir = os.getcwd() + os.sep + "output"
        if not os.path.exists(self._output_dir):
            os.mkdir(self._output_dir)
        assert os.path.isdir(self._output_dir),\
            self._output_dir + " is not a directory."

    def _create_filenames(self):
        filebase = self.step.scenario.feature.name + "-" +\
                   self.step.scenario.name + "-" + self.process_name
        self.stderr_filename = self._output_dir + os.sep +\
                               self.mangle_filename(filebase, "stderr")
        self.stdout_filename = self._output_dir + os.sep +\
                               self.mangle_filename(filebase, "stdout")

    def stop_process(self):
        if self.process is not None:
            self.process.terminate()
            self.process.wait()
        self.process = None
        if self.remove_files_on_exit:
            self._remove_files()

    def _remove_files(self):
        os.remove(self.stderr_filename)
        os.remove(self.stdout_filename)

    def _wait_for_output_str(self, filename, running_file, strings, only_new):
        if not only_new:
            full_file = open(filename, "r")
            for line in full_file:
                for string in strings:
                    if line.find(string) != -1:
                        full_file.close()
                        return string
        wait_count = 0
        while wait_count < OUTPUT_WAIT_MAX_INTERVALS:
            where = running_file.tell()
            line = running_file.readline()
            if line:
                for string in strings:
                    if line.find(string) != -1:
                        return string
            else:
                wait_count += 1
                time.sleep(OUTPUT_WAIT_INTERVAL)
                running_file.seek(where)
        assert False, "Timeout waiting for process output: " + str(strings)

    def wait_for_stderr_str(self, strings, only_new = True):
        return self._wait_for_output_str(self.stderr_filename, self.stderr,
                                         strings, only_new)

    def wait_for_stdout_str(self, strings, only_new = True):
        return self._wait_for_output_str(self.stdout_filename, self.stdout,
                                         strings, only_new)

# Container class for a number of running processes
# i.e. servers like bind10, etc
# one-shot programs like dig or bindctl are started and closed separately
class RunningProcesses:
    def __init__(self):
        self.processes = {}
    
    def add_process(self, step, process_name, args):
        assert process_name not in self.processes,\
            "Process " + name + " already running"
        self.processes[process_name] = RunningProcess(step, process_name, args)

    def get_process(self, process_name):
        assert process_name in self.processes,\
            "Process " + name + " unknown"
        return self.processes[process_name]

    def stop_process(self, process_name):
        assert process_name in self.processes,\
            "Process " + name + " unknown"
        self.processes[process_name].stop_process()
        del self.processes[process_name]
        
    def stop_all_processes(self):
        for process in self.processes.values():
            process.stop_process()
    
    def keep_files(self):
        for process in self.processes.values():
            process.remove_files_on_exit = False

    def wait_for_stderr_str(self, process_name, strings, only_new = True):
        """Wait for any of the given strings in the given processes stderr 
        output. If only_new is True, it will only look at the lines that are 
        printed to stderr since the last time this method was called. If 
        False, it will also look at the previously printed lines. This will 
        block until one of the strings is found. TODO: we may want to put in 
        a timeout for this... Returns the string that is found"""
        assert process_name in self.processes,\
           "Process " + process_name + " unknown"
        return self.processes[process_name].wait_for_stderr_str(strings,
                                                                only_new)

    def wait_for_stdout_str(self, process_name, strings, only_new = True):
        """Wait for any of the given strings in the given processes stderr 
        output. If only_new is True, it will only look at the lines that are 
        printed to stderr since the last time this method was called. If 
        False, it will also look at the previously printed lines. This will 
        block until one of the strings is found. TODO: we may want to put in 
        a timeout for this... Returns the string that is found"""
        assert process_name in self.processes,\
           "Process " + process_name + " unknown"
        return self.processes[process_name].wait_for_stdout_str(strings,
                                                                only_new)

@before.each_scenario
def initialize(scenario):
    # Keep track of running processes
    world.processes = RunningProcesses()

    # Convenience variable to access the last query result from querying.py
    world.last_query_result = None

    # Some tests can modify the settings. If the tests fail half-way, or
    # don't clean up, this can leave configurations or data in a bad state,
    # so we copy them from originals before each scenario
    for item in copylist:
        shutil.copy(item[0], item[1])

    for item in removelist:
        if os.path.exists(item):
            os.remove(item)

@after.each_scenario
def cleanup(scenario):
    # Keep output files if the scenario failed
    if not scenario.passed:
        world.processes.keep_files()
    # Stop any running processes we may have had around
    world.processes.stop_all_processes()