wuttjamaican/tests/test_batch.py

# -*- coding: utf-8; -*-

import os
from unittest.mock import patch

from wuttjamaican import batch as mod

try:
    import sqlalchemy as sa
    from wuttjamaican.db import model
    from wuttjamaican.testing import DataTestCase
except ImportError:
    pass
else:

    class MockBatch(model.BatchMixin, model.Base):
        __tablename__ = 'testing_batch_mock'

    class MockBatchRow(model.BatchRowMixin, model.Base):
        __tablename__ = 'testing_batch_mock_row'
        __batch_class__ = MockBatch

    class MockBatchHandler(mod.BatchHandler):
        model_class = MockBatch

    class TestBatchHandler(DataTestCase):

        def make_handler(self, **kwargs):
            return MockBatchHandler(self.config, **kwargs)

        def test_model_class(self):
            handler = mod.BatchHandler(self.config)
            self.assertRaises(NotImplementedError, getattr, handler, 'model_class')

        def test_batch_type(self):
            with patch.object(mod.BatchHandler, 'model_class', new=MockBatch):
                handler = mod.BatchHandler(self.config)
                self.assertEqual(handler.batch_type, 'testing_batch_mock')

        def test_make_batch(self):
            handler = self.make_handler()
            batch = handler.make_batch(self.session)
            self.assertIsInstance(batch, MockBatch)

        def test_consume_batch_id(self):
            handler = self.make_handler()

            first = handler.consume_batch_id(self.session)
            second = handler.consume_batch_id(self.session)
            self.assertEqual(second, first + 1)

            third = handler.consume_batch_id(self.session, as_str=True)
            self.assertEqual(third, f'{first + 2:08d}')

        def test_get_data_path(self):
            model = self.app.model
            user = model.User(username='barney')
            self.session.add(user)

            with patch.object(mod.BatchHandler, 'model_class', new=MockBatch):
                handler = self.make_handler()

                # root storage (default)
                with patch.object(self.app, 'get_appdir', return_value=self.tempdir):
                    path = handler.get_data_path()
                    self.assertEqual(path, os.path.join(self.tempdir, 'data', 'batch', 'testing_batch_mock'))

                # root storage (configured)
                self.config.setdefault('wutta.batch.storage_path', self.tempdir)
                path = handler.get_data_path()
                self.assertEqual(path, os.path.join(self.tempdir, 'testing_batch_mock'))

                batch = handler.make_batch(self.session, created_by=user)
                self.session.add(batch)
                self.session.flush()

                # batch-specific
                path = handler.get_data_path(batch)
                uuid = batch.uuid.hex
                final = os.path.join(uuid[-2:], uuid[:-2])
                self.assertEqual(path, os.path.join(self.tempdir, 'testing_batch_mock', final))

                # with filename
                path = handler.get_data_path(batch, 'input.csv')
                self.assertEqual(path, os.path.join(self.tempdir, 'testing_batch_mock', final, 'input.csv'))

                # makedirs
                path = handler.get_data_path(batch)
                self.assertFalse(os.path.exists(path))
                path = handler.get_data_path(batch, makedirs=True)
                self.assertTrue(os.path.exists(path))

        def test_should_populate(self):
            handler = self.make_handler()
            batch = handler.make_batch(self.session)
            self.assertFalse(handler.should_populate(batch))

        def test_do_populate(self):
            handler = self.make_handler()
            batch = handler.make_batch(self.session)
            # nb. coverage only; tests nothing
            handler.do_populate(batch)

        def test_make_row(self):
            handler = self.make_handler()
            row = handler.make_row()
            self.assertIsInstance(row, MockBatchRow)

        def test_add_row(self):
            handler = self.make_handler()
            batch = handler.make_batch(self.session)
            self.session.add(batch)
            row = handler.make_row()
            self.assertIsNone(batch.row_count)
            handler.add_row(batch, row)
            self.assertEqual(batch.row_count, 1)

        def test_do_execute(self):
            model = self.app.model
            user = model.User(username='barney')
            self.session.add(user)

            handler = self.make_handler()
            batch = handler.make_batch(self.session, created_by=user)
            self.session.add(batch)
            self.session.flush()

            # error if execution not allowed
            with patch.object(handler, 'why_not_execute', return_value="bad batch"):
                self.assertRaises(RuntimeError, handler.do_execute, batch, user)

            # nb. coverage only; tests nothing
            self.assertIsNone(batch.executed)
            self.assertIsNone(batch.executed_by)
            handler.do_execute(batch, user)
            self.assertIsNotNone(batch.executed)
            self.assertIs(batch.executed_by, user)

            # error if execution already happened
            self.assertRaises(ValueError, handler.do_execute, batch, user)

        def test_do_delete(self):
            model = self.app.model
            handler = self.make_handler()

            user = model.User(username='barney')
            self.session.add(user)

            # simple delete
            batch = handler.make_batch(self.session, created_by=user)
            self.session.add(batch)
            self.session.flush()
            self.assertEqual(self.session.query(MockBatch).count(), 1)
            handler.do_delete(batch, user)
            self.assertEqual(self.session.query(MockBatch).count(), 0)

            # delete w/ rows
            batch = handler.make_batch(self.session, created_by=user)
            self.session.add(batch)
            for i in range(5):
                row = handler.make_row()
                handler.add_row(batch, row)
            self.session.flush()
            self.assertEqual(self.session.query(MockBatch).count(), 1)
            handler.do_delete(batch, user)
            self.assertEqual(self.session.query(MockBatch).count(), 0)

            # delete w/ files
            self.config.setdefault('wutta.batch.storage_path', self.tempdir)
            batch = handler.make_batch(self.session, created_by=user)
            self.session.add(batch)
            self.session.flush()
            path = handler.get_data_path(batch, 'data.txt', makedirs=True)
            with open(path, 'wt') as f:
                f.write('foo=bar')
            self.assertEqual(self.session.query(MockBatch).count(), 1)
            path = handler.get_data_path(batch)
            self.assertTrue(os.path.exists(path))
            handler.do_delete(batch, user)
            self.assertEqual(self.session.query(MockBatch).count(), 0)
            self.assertFalse(os.path.exists(path))

            # delete w/ files (dry-run)
            self.config.setdefault('wutta.batch.storage_path', self.tempdir)
            batch = handler.make_batch(self.session, created_by=user)
            self.session.add(batch)
            self.session.flush()
            path = handler.get_data_path(batch, 'data.txt', makedirs=True)
            with open(path, 'wt') as f:
                f.write('foo=bar')
            self.assertEqual(self.session.query(MockBatch).count(), 1)
            path = handler.get_data_path(batch)
            self.assertTrue(os.path.exists(path))
            handler.do_delete(batch, user, dry_run=True)
            # nb. batch appears missing from session even in dry-run
            self.assertEqual(self.session.query(MockBatch).count(), 0)
            # nb. but its files remain intact
            self.assertTrue(os.path.exists(path))