mrjob's unit test example

mr_counting_job.py

from mrjob.job  import MRJob
from mrjob.step import MRStep

class MRCountingJob(MRJob):

    def __init__(self, *args, **kwargs):
        super(MRCountingJob, self).__init__(*args, **kwargs)

    def steps(self):
        # 3 steps so we can check behavior of counters for multiple steps
        return [self.mr(mapper=self.mapper),
                self.mr(mapper=self.mapper),
                self.mr(mapper=self.mapper)]

    def mapper(self, _, value):
        self.increment_counter('group', 'counter_name', 1)
        yield _, value

if __name__ == '__main__':
    MRCountingJob.run()

test_counters.py

from __future__ import with_statement

try:
    import unittest2 as unittest
except ImportError:
    import unittest

from StringIO import StringIO
from mr_counting_job import MRCountingJob

class CounterTestCase(unittest.TestCase):

    def test_counters(self):
        stdin = StringIO('foo\nbar\n')

        mr_job = MRCountingJob(['--no-conf', '-'])
        mr_job.sandbox(stdin=stdin)

        with mr_job.make_runner() as runner:
            runner.run()

            self.assertEqual(runner.counters(),
                             [{'group': {'counter_name': 2}},
                              {'group': {'counter_name': 2}},
                              {'group': {'counter_name': 2}}])

if __name__ == '__main__':
    unittest.main()