Repository: spotify/luigi
File: test/contrib/postgres_test.py (Python)

Code before:
    def test_bulk_complete(self, mock_connect):
        mock_cursor = MockPostgresCursor([
            'DummyPostgresQuery_2015_01_03_838e32a989'
        ])
        mock_connect.return_value.cursor.return_value = mock_cursor

        task = RangeDaily(of=DummyPostgresQuery,
                          start=datetime.date(2015, 1, 2),
                          now=datetime_to_epoch(datetime.datetime(2015, 1, 7)))
        actual = [t.task_id for t in task.requires()]

        self.assertEqual(actual, [
            'DummyPostgresQuery_2015_01_02_3a0ec498ed',
            'DummyPostgresQuery_2015_01_04_9c1d42ff62',
            'DummyPostgresQuery_2015_01_05_0f90e52357',
            'DummyPostgresQuery_2015_01_06_f91a47ec40',
        ])
        self.assertFalse(task.complete())

    def test_override_port(self):
        output = DummyPostgresQueryWithPort(date=datetime.datetime(1991, 3, 24)).output()
        self.assertEquals(output.port, 1234)


@attr('postgres')
class TestCopyToTableWithMetaColumns(unittest.TestCase):
    @mock.patch("luigi.contrib.postgres.CopyToTable.enable_metadata_columns", new_callable=mock.PropertyMock, return_value=True)
    @mock.patch("luigi.contrib.postgres.CopyToTable._add_metadata_columns")
    @mock.patch("luigi.contrib.postgres.CopyToTable.post_copy_metacolumns")
    @mock.patch("luigi.contrib.postgres.CopyToTable.rows", return_value=['row1', 'row2'])
    @mock.patch("luigi.contrib.postgres.PostgresTarget")
    @mock.patch('psycopg2.connect')
    def test_copy_with_metadata_columns_enabled(self,
                                                mock_connect,
                                                mock_redshift_target,
                                                mock_rows,
                                                mock_add_columns,
                                                mock_update_columns,
                                                mock_metadata_columns_enabled):

        task = DummyPostgresImporter(date=datetime.datetime(1991, 3, 24))

        mock_cursor = MockPostgresCursor([task.task_id])
        mock_connect.return_value.cursor.return_value = mock_cursor

        task = DummyPostgresImporter(date=datetime.datetime(1991, 3, 24))
        task.run()

        self.assertTrue(mock_add_columns.called)
        self.assertTrue(mock_update_columns.called)


Review comment (suggestion):
Could we also add a test for the `host:port` parsing? I'm not seeing one here

Code after:
class PostgresQueryTest(unittest.TestCase):
    maxDiff = None

    @mock.patch('psycopg2.connect')
    def test_bulk_complete(self, mock_connect):
        mock_cursor = MockPostgresCursor([
            'DummyPostgresQuery_2015_01_03_838e32a989'
        ])
        mock_connect.return_value.cursor.return_value = mock_cursor

        task = RangeDaily(of=DummyPostgresQuery,
                          start=datetime.date(2015, 1, 2),
                          now=datetime_to_epoch(datetime.datetime(2015, 1, 7)))
        actual = [t.task_id for t in task.requires()]

        self.assertEqual(actual, [
            'DummyPostgresQuery_2015_01_02_3a0ec498ed',
            'DummyPostgresQuery_2015_01_04_9c1d42ff62',
            'DummyPostgresQuery_2015_01_05_0f90e52357',
            'DummyPostgresQuery_2015_01_06_f91a47ec40',
        ])
        self.assertFalse(task.complete())

    def test_override_port(self):
        output = DummyPostgresQueryWithPort(date=datetime.datetime(1991, 3, 24)).output()
        self.assertEquals(output.port, 1234)

    def test_port_encoded_in_host(self):
        output = DummyPostgresQueryWithPortEncodedInHost(date=datetime.datetime(1991, 3, 24)).output()
        self.assertEquals(output.port, '1234')


@attr('postgres')
class TestCopyToTableWithMetaColumns(unittest.TestCase):
    @mock.patch("luigi.contrib.postgres.CopyToTable.enable_metadata_columns", new_callable=mock.PropertyMock, return_value=True)
    @mock.patch("luigi.contrib.postgres.CopyToTable._add_metadata_columns")
    @mock.patch("luigi.contrib.postgres.CopyToTable.post_copy_metacolumns")
    @mock.patch("luigi.contrib.postgres.CopyToTable.rows", return_value=['row1', 'row2'])
    @mock.patch("luigi.contrib.postgres.PostgresTarget")
    @mock.patch('psycopg2.connect')
    def test_copy_with_metadata_columns_enabled(self,
                                                mock_connect,
                                                mock_redshift_target,
                                                mock_rows,
                                                mock_add_columns,
                                                mock_update_columns,
                                                mock_metadata_columns_enabled):

        task = DummyPostgresImporter(date=datetime.datetime(1991, 3, 24))

        mock_cursor = MockPostgresCursor([task.task_id])