nilmdb/tests/test_iteratorizer.py

import nilmdb
from nilmdb.utils.printf import *

import nose
from nose.tools import *
from nose.tools import assert_raises
import threading
import time

from testutil.helpers import *

def func_with_callback(a, b, callback):
    callback(a)
    callback(b)
    callback(a+b)
    return "return value"

class TestIteratorizer(object):
    def test(self):

        # First try it with a normal callback
        self.result = ""
        def cb(x):
            self.result += str(x)
        func_with_callback(1, 2, cb)
        eq_(self.result, "123")

        # Now make it an iterator
        result = ""
        f = lambda x: func_with_callback(1, 2, x)
        with nilmdb.utils.Iteratorizer(f) as it:
            for i in it:
                result += str(i)
        eq_(result, "123")
        eq_(it.retval, "return value")

        # Make sure things work when an exception occurs
        result = ""
        with nilmdb.utils.Iteratorizer(
            lambda x: func_with_callback(1, "a", x)) as it:
            with assert_raises(TypeError) as e:
                for i in it:
                    result += str(i)
        eq_(result, "1a")

        # Now try to trigger the case where we stop iterating
        # mid-generator, and expect the iteratorizer to clean up after
        # itself.  This doesn't have a particular result in the test,
        # but gains coverage.
        def foo():
            with nilmdb.utils.Iteratorizer(f) as it:
                it.next()
        foo()
        eq_(it.retval, None)

        # Do the same thing when the curl hack is applied
        def foo():
            with nilmdb.utils.Iteratorizer(f, curl_hack = True) as it:
                it.next()
        foo()
        eq_(it.retval, None)