circleci_join.py

#!/usr/bin/env python
import os
import sys
import time

import psutil

project_root = os.getcwd()
sys.path.insert(0, project_root)

from utils.shell import TmuxSession

session_name = sys.argv[1]
session = TmuxSession(session_name)


def process_terminated(process):
    print('process_terminated')
    session_terminated(process.returncode)


def session_terminated(retcode=None):
    print('session_terminated({})'.format(retcode))
    session.print_history()
    if retcode is None:
        retcode = session.fetch_return_code()
    # session.quit()
    session.send('Would be quitting...')
    exit(retcode)


if not session.exists():
    print('No such session: {}'.format(session_name))
elif session.idle:
    session_terminated()
else:
    print('Waiting for process to finish: {}, {}'.format(time.time(), session.child_pid))
    alive = [session.child_process]
    while alive:
        gone, alive = psutil.wait_procs([session.child_process], timeout=60, callback=process_terminated)
        print('Timed out after 60, still alive: {}'.format(alive))

circleci_run_bg.py

#!/usr/bin/env python
import argparse
import sys

import os

project_root = os.getcwd()
sys.path.insert(0, project_root)

from utils.shell import TmuxSession

parser = argparse.ArgumentParser()
parser.add_argument('--venv', action='store_true', default=False)
parser.add_argument('session_name')
parser.add_argument('command')
args = parser.parse_args()

tmux = TmuxSession(args.session_name)
if args.venv:
    venv_path = os.environ['VIRTUAL_ENV']
    tmux.run('. {}'.format(os.path.join(venv_path, 'bin/activate')))
tmux.run(args.command)
print(tmux.child_pid)

utils_shell.py

class TmuxSession(object):

    def __init__(self, session_name):
        self._session_name = session_name
        self._pid = None
        self._process = None
        self._started = False

    def _execute_no_session(self, *args):
        """
        Run tmux command that is not necessarily bound to this session
        """
        return run_non_interactive(['tmux'] + list(args))

    def _execute(self, cmdname, *args, **kwargs):
        """
        Run tmux command on this session, creating it first if it doesn't already exist
        """
        if not self._started and not self.exists():
            retcode, out, err = self._execute_no_session('new', '-d', '-s', self._session_name, '-x', 272)
            assert retcode == 0, ('Could not create session: {}'.format(retcode), out, err)
        self._started = True

        expected_ret_code = kwargs.pop('expected_ret_code', 0)
        assert not kwargs, kwargs

        args = [cmdname, '-t', self._session_name] + list(args)

        retcode, out, err = self._execute_no_session(*args)
        assert retcode == expected_ret_code, \
            ('Expected retcode: {}, got: {}'.format(expected_ret_code, retcode), out, err)
        return retcode, out, err

    def exists(self):
        retcode, out, err = self._execute_no_session('has-session', '-t', self._session_name)
        return retcode == 0

    def send(self, *args):
        return self._execute('send', *args)

    def shell_execute(self, text, outfile=None):
        args = [text]
        if outfile is not None:
            args += [' > {}'.format(outfile)]
        args.append('ENTER')
        self.send(*args)

    def quit(self):
        self.send('^D')

    @property
    def process(self):
        if not self._process:
            self._process = psutil.Process(self.pid)
        return self._process

    @property
    def child_process(self):
        children = self.process.children()
        return children[0] if children else None

    @property
    def child_pid(self):
        p = self.child_process
        return p.pid if p else None

    def fetch_return_code(self):
        assert self.idle
        outfd, outfname = mkstemp(dir='/tmp')
        self.shell_execute('echo $?', outfile=outfname)
        # Give the above command some time to finish
        time.sleep(.5)
        with open(outfname) as f:
            retcode = int(f.read())
        os.unlink(outfname)
        return retcode

    @property
    def idle(self):
        return not self.child_process

    @property
    def pid(self):
        if self._pid is None:
            retcode, out, err = self._execute('list-panes', '-F', '#{pane_pid}')
            assert retcode == 0, (retcode, out, err)
            lines = out.strip().splitlines()
            assert len(lines) == 1, 'Expected 1 line, got {}:\n{}'.format(len(lines), lines)
            self._pid = int(lines[0])
        return self._pid

    def run(self, cmd):
        self.shell_execute(cmd)
        return self.child_pid

    def print_history(self):
        subprocess.Popen(['tmux', 'capture-pane', '-t', self._session_name, '-e', '-p', '-S', '-100000']).communicate()