saucebrush/tests/emitters.py

from contextlib import closing
from io import StringIO
import os
import unittest

from saucebrush.emitters import (
    DebugEmitter,
    CSVEmitter,
    CountEmitter,
    SqliteEmitter,
    SqlDumpEmitter,
)


class EmitterTestCase(unittest.TestCase):
    def test_debug_emitter(self):
        with closing(StringIO()) as output:
            de = DebugEmitter(output)
            list(de.attach([1, 2, 3]))
            self.assertEqual(output.getvalue(), "1\n2\n3\n")

    def test_count_emitter(self):

        # values for test
        values = [
            1,
            2,
            3,
            4,
            5,
            6,
            7,
            8,
            9,
            10,
            11,
            12,
            13,
            14,
            15,
            16,
            17,
            18,
            19,
            20,
            21,
            22,
        ]

        with closing(StringIO()) as output:

            # test without of parameter
            ce = CountEmitter(every=10, outfile=output, format="%(count)s records\n")
            list(ce.attach(values))
            self.assertEqual(output.getvalue(), "10 records\n20 records\n")
            ce.done()
            self.assertEqual(output.getvalue(), "10 records\n20 records\n22 records\n")

        with closing(StringIO()) as output:

            # test with of parameter
            ce = CountEmitter(every=10, outfile=output, of=len(values))
            list(ce.attach(values))
            self.assertEqual(output.getvalue(), "10 of 22\n20 of 22\n")
            ce.done()
            self.assertEqual(output.getvalue(), "10 of 22\n20 of 22\n22 of 22\n")

    def test_csv_emitter(self):

        try:
            import cStringIO  # if Python 2.x then use old cStringIO

            io = cStringIO.StringIO()
        except Exception:
            io = StringIO()  # if Python 3.x then use StringIO

        with closing(io) as output:
            ce = CSVEmitter(output, ("x", "y", "z"))
            list(ce.attach([{"x": 1, "y": 2, "z": 3}, {"x": 5, "y": 5, "z": 5}]))
            self.assertEqual(output.getvalue(), "x,y,z\r\n1,2,3\r\n5,5,5\r\n")

    def test_sqlite_emitter(self):

        import sqlite3
        import tempfile

        with closing(tempfile.NamedTemporaryFile(suffix=".db")) as f:
            db_path = f.name

        sle = SqliteEmitter(db_path, "testtable", fieldnames=("a", "b", "c"))
        list(sle.attach([{"a": "1", "b": "2", "c": "3"}]))
        sle.done()

        with closing(sqlite3.connect(db_path)) as conn:
            cur = conn.cursor()
            cur.execute("""SELECT a, b, c FROM testtable""")
            results = cur.fetchall()

        os.unlink(db_path)

        self.assertEqual(results, [("1", "2", "3")])

    def test_sql_dump_emitter(self):

        with closing(StringIO()) as bffr:

            sde = SqlDumpEmitter(bffr, "testtable", ("a", "b"))
            list(sde.attach([{"a": 1, "b": "2"}]))
            sde.done()

            self.assertEqual(
                bffr.getvalue(), "INSERT INTO `testtable` (`a`,`b`) VALUES (1,'2');\n"
            )


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