tests.py

import time
import genquery
import oldquery
from collections import OrderedDict

# Call as `irule run_test null ruleExecOut`
#
# genquery.py should be genquery from the fork: https://github.com/cjsmeele/irods_rule_engine_plugin_python/blob/genquery-extensions/genquery.py
# oldquery.py should be the current genquery: https://github.com/irods/irods_rule_engine_plugin_python/blob/master/genquery.py

class Tracker(object):
    """Count msi calls (callback wrapper)"""
    def __init__(self, callback):
        self._callback = callback
        self._stats = {}
    def __setattr__(self, k, v):
        if k.startswith('_'):
            super(Tracker, self).__setattr__(k, v)
        else:
            self._callback.k = v
    def __getattr__(self, k):
        if k.startswith('_'):
            return super(Tracker, self).__getattr__(k)
        elif k.startswith('msi'):
            if k in self._stats:
                self._stats[k] += 1
            else:
                self._stats[k] = 1
        return getattr(self._callback, k)
    def stats(self):
        return self._stats


def f(callback, q, abort):
    x = time.time()
    if abort:
        # Get only one row.
        for z in q.row_iterator('COLL_NAME', '', q.AS_LIST, callback):
            break
    else:
        # Consume all rows.
        len(list(q.row_iterator('COLL_NAME', '', q.AS_LIST, callback)))

    y = time.time()
    return y - x

def measure(callback, q, n, abort):
    """Run a query n times via the q query module"""
    callback = Tracker(callback)
    ts = []
    for i in range(n):
        x   = f(callback, q, abort)
        ts += [x]
    return OrderedDict([('avg', sum(ts)/len(ts)),
                        ('max', max(ts)),
                        ('min', min(ts)),
                        ('msi_stats', callback.stats())])

def run_test(rule_args, callback, rei):
    def run_1(abort):
        # (throw away results of one test to warm-up)
        measure(callback, oldquery, 1, abort)
        old = measure(callback, oldquery, 10, abort)
        new = measure(callback, genquery, 10, abort)

        return OrderedDict([('old', old),
                            ('new', new),
                            ('avg_improvement_ratio', old['avg']/new['avg'])])

    import json
    callback.writeLine('stdout', json.dumps(
                                    OrderedDict([('with_completed_query', run_1(False)),
                                                 ('with_query_aborted_after_1_result', run_1(True))]), indent=4))