Source code for idaes.util.testutil

##############################################################################
# Institute for the Design of Advanced Energy Systems Process Systems
# Engineering Framework (IDAES PSE Framework) Copyright (c) 2018, by the
# software owners: The Regents of the University of California, through
# Lawrence Berkeley National Laboratory,  National Technology & Engineering
# Solutions of Sandia, LLC, Carnegie Mellon University, West Virginia
# University Research Corporation, et al. All rights reserved.
# 
# Please see the files COPYRIGHT.txt and LICENSE.txt for full copyright and
# license information, respectively. Both files are also available online
# at the URL "https://github.com/IDAES/idaes".
##############################################################################
"""
Utility functions, etc. for testing
"""
# stdlib
import glob
import os
import re
import signal
import subprocess
import time
import warnings
# third-party
import nbformat
from nbconvert.preprocessors import ExecutePreprocessor
import psutil
# idaes
import idaes


[docs]class RunnableExample(object): """Encapsulate a runnable example script. Example usage (e.g. in a test): ex = RunnableExample("dmf/resource_example.py") rcode = ex.run(timeout=5) assert rcode == 0, "dmf/resource_example.py failed" """ ansi_escape = re.compile(r'\x1B\[[0-?]*[ -/]*[@-~]') def __init__(self, path): """Create new example with given script path. Args: path(str): Relative (to 'examples/' dir) or absolute path to the script. """ self._path = path self._reset() def _reset(self): self._status = -1 self._out, self._err = '', ''
[docs] def run(self, args=None, timeout=15): """Run the example. * Files ending in '.py' are run with the Python interpreter. * Files ending in '.ipynb' (Jupyter notebooks) are executed using Jupyter nbconvert's pre-processing API. * All other files are simply executed. Args: args (list[str]): Arguments passed to scripts (except notebooks) timeout (int): Time, in seconds, to wait for it to run. Returns: Numeric status code """ self._reset() args = [] if args is None else args script = self._get_script_path() if script.endswith('.ipynb'): return self._run_notebook(script, timeout) if script.endswith('.py'): cmd = ['python', script] + args else: cmd = [script] + args # print('@@ cmd={}'.format(cmd)) proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) t0, status = time.time(), None while status is None and time.time() - t0 <= timeout: time.sleep(1) status = proc.poll() if status is None: dead = self._terminate(proc) if not dead: warnings.warn('Unable to terminate process {}' .format(proc.pid)) return -1 if status != 0: self._out = proc.stdout.read() self._err = proc.stderr.read() return status
def _run_notebook(self, nbfile, timeout): """Special case for running a Jupyter Notebook. See: http://nbconvert.readthedocs.io/en/latest/execute_api.html """ status = 0 with open(nbfile) as f: nb = nbformat.read(f, as_version=4) ep = ExecutePreprocessor(timeout=timeout + 1) run_path = os.path.dirname(nbfile) try: ep.preprocess(nb, {'metadata': {'path': run_path}}) except Exception as err: self._err = self.ansi_escape.sub('', str(err)) self._out = 'Error executing the notebook "{}"'.format(nbfile) status = 1 return status @property def output(self): return self._out @property def error(self): return self._err @staticmethod def _terminate(proc): pid = proc.pid result = True os.kill(pid, signal.SIGTERM) time.sleep(5) if psutil.pid_exists(pid): os.kill(pid, signal.SIGKILL) time.sleep(5) if psutil.pid_exists(pid): result = False return result def _get_script_path(self): if os.path.isabs(self._path): return self._path return os.path.join(examples_path(), self._path)
[docs]def examples_path(): """Return path to examples directory. """ return os.path.join(os.path.dirname(idaes.__path__[0]), 'examples')
[docs]def find_examples(subdir): root = os.path.join(examples_path(), subdir) examples = glob.glob(os.path.join(root, '*_example.py')) return examples