timewarrior/test/basetest/timew.py

# -*- coding: utf-8 -*-

import atexit
import datetime
import json
import os
import shlex
import shutil
import tempfile
import unittest

from .compat import STRING_TYPE
from .exceptions import CommandError
from .utils import run_cmd_wait, run_cmd_wait_nofail, which, timew_binary_location, DEFAULT_EXTENSION_PATH


class Timew(object):
    """Manage a Timewarrior instance

    A temporary folder is used as data store of timewarrior.

    A timew client should not be used after being destroyed.
    """
    DEFAULT_TIMEW = timew_binary_location()

    def __init__(self, timew=DEFAULT_TIMEW):
        """Initialize a timewarrior (client).
        The program runs in a temporary folder.

        :arg timew: Timewarrior binary to use as client (defaults: timew in PATH)
        """
        self.timew = timew

        # Used to specify what command to launch (and to inject faketime)
        self._command = [self.timew]

        # Configuration of the isolated environment
        self._original_pwd = os.getcwd()
        self.datadir = tempfile.mkdtemp(prefix="timew_")
        self.timewrc = os.path.join (self.datadir, 'timewarrior.cfg')
        self.extdir = os.path.join(self.datadir, 'extensions')

        # Ensure any instance is properly destroyed at session end
        atexit.register(lambda: self.destroy())

        self.reset_env()

    def add_default_extension(self, filename):
        """Add default extension to current instance"""
        if not os.path.isdir(self.extdir):
            os.mkdir(self.extdir)

        extfile = os.path.join(self.extdir, filename)
        if os.path.isfile(extfile):
            raise "{} already exists".format(extfile)

        shutil.copy(os.path.join(DEFAULT_EXTENSION_PATH, filename), extfile)

    def __repr__(self):
        txt = super(Timew, self).__repr__()
        return "{0} running from {1}>".format(txt[:-1], self.datadir)

    def __call__(self, *args, **kwargs):
        """aka t = Timew() ; t() which is now an alias to t.runSuccess()"""
        return self.runSuccess(*args, **kwargs)

    def reset_env(self):
        """Set a new environment derived from the one used to launch the test"""
        # Copy all env variables to avoid clashing subprocess environments
        self.env = os.environ.copy()

        # As well as TIMEWARRIORDB
        self.env["TIMEWARRIORDB"] = self.datadir

    def config(self, var, value):
        """Run setup `var` as `value` in timew config"""
        # Add -- to avoid misinterpretation of - in things like UUIDs
        # TODO Revisit this decision as timew matures.
        # cmd = (self.timew, "config", "--", var, value)
        cmd = (self.timew, ":yes", "config", var, value)
        return run_cmd_wait(cmd, env=self.env, input="y\n")

    def configure_exclusions(self, start, end):
        if isinstance(start, datetime.time) and isinstance(end, datetime.time):
            if start > end:
                exclusion = "<{:%H:%M:%S} >{:%H:%M:%S}".format(end, start)
            else:
                exclusion = "{:%H:%M:%S}-{:%H:%M:%S}".format(start, end)
        else:
            exclusion = "{}-{}".format(start, end)

        self.config("exclusions.monday", exclusion)
        self.config("exclusions.tuesday", exclusion)
        self.config("exclusions.wednesday", exclusion)
        self.config("exclusions.thursday", exclusion)
        self.config("exclusions.friday", exclusion)
        self.config("exclusions.saturday", exclusion)
        self.config("exclusions.sunday", exclusion)

    def del_config(self, var):
        """Remove `var` from timew config"""
        cmd = (self.timew, ":yes", "config", var)
        return run_cmd_wait(cmd, env=self.env, input="y\n")

    @property
    def timewrc_content(self):
        """Returns the contents of the timewrc file."""

        with open(self.timewrc, "r") as f:
            return f.readlines()

    def export(self, export_filter=None):
        """Run "timew export", return JSON array of exported intervals."""
        if export_filter is None:
            export_filter = ""

        code, out, err = self.runSuccess("{0} export".format(export_filter))

        return json.loads(out)

    @staticmethod
    def _split_string_args_if_string(args):
        """Helper function to parse and split into arguments a single string
        argument. The string is literally the same as if written in the shell.
        """
        # Enable nicer-looking calls by allowing plain strings
        if isinstance(args, STRING_TYPE):
            args = shlex.split(args)

        return args

    def runSuccess(self, args="", input=None, merge_streams=False, timeout=5):
        """Invoke timew with given arguments and fail if exit code != 0

        Use runError if you want exit_code to be tested automatically and
        *not* fail if program finishes abnormally.

        If you wish to pass instructions to timew such as confirmations or other
        input via stdin, you can do so by providing a input string.
        Such as input="y\ny\n".

        If merge_streams=True stdout and stderr will be merged into stdout.

        timeout = number of seconds the test will wait for every timew call.
        Defaults to 1 second if not specified. Unit is seconds.

        Returns (exit_code, stdout, stderr) if merge_streams=False
                (exit_code, output) if merge_streams=True
        """
        # Create a copy of the command
        command = self._command[:]

        args = self._split_string_args_if_string(args)
        command.extend(args)

        output = run_cmd_wait_nofail(command, input,
                                     merge_streams=merge_streams,
                                     env=self.env,
                                     timeout=timeout)

        if output[0] != 0:
            raise CommandError(command, *output)

        return output

    def runError(self, args=(), input=None, merge_streams=False, timeout=5):
        """Invoke timew with given arguments and fail if exit code == 0

        Use runSuccess if you want exit_code to be tested automatically and
        *fail* if program finishes abnormally.

        If you wish to pass instructions to timew such as confirmations or other
        input via stdin, you can do so by providing a input string.
        Such as input="y\ny\n".

        If merge_streams=True stdout and stderr will be merged into stdout.

        timeout = number of seconds the test will wait for every timew call.
        Defaults to 1 second if not specified. Unit is seconds.

        Returns (exit_code, stdout, stderr) if merge_streams=False
                (exit_code, output) if merge_streams=True
        """
        # Create a copy of the command
        command = self._command[:]

        args = self._split_string_args_if_string(args)
        command.extend(args)

        output = run_cmd_wait_nofail(command, input,
                                     merge_streams=merge_streams,
                                     env=self.env,
                                     timeout=timeout)

        # output[0] is the exit code
        if output[0] == 0 or output[0] is None:
            raise CommandError(command, *output)

        return output

    def destroy(self):
        """Cleanup the data folder and release server port for other instances"""
        try:
            shutil.rmtree(self.datadir)
        except OSError as e:
            if e.errno == 2:
                # Directory no longer exists
                pass
            else:
                raise

        # Prevent future reuse of this instance
        self.runSuccess = self.__destroyed
        self.runError = self.__destroyed

        # self.destroy will get called when the python session closes.
        # If self.destroy was already called, turn the action into a noop
        self.destroy = lambda: None

    def __destroyed(self, *args, **kwargs):
        raise AttributeError("Program instance has been destroyed. "
                             "Create a new instance if you need a new client.")

    def faketime(self, faketime=None):
        """Set a faketime using libfaketime that will affect the following
        command calls.

        If faketime is None, faketime settings will be disabled.
        """
        cmd = which("faketime")
        if cmd is None:
            raise unittest.SkipTest("libfaketime/faketime is not installed")

        if self._command[0] == cmd:
            self._command = self._command[3:]

        if faketime is not None:
            # Use advanced time format
            self._command = [cmd, "-f", faketime] + self._command

# vim: ai sts=4 et sw=4