Beispiel #1
0
    def test_should_support_db_backend(self, mock_session):
        with _prepare_app():
            mock_backend = DatabaseBackend(app=celery_executor.app,
                                           url="sqlite3://")

            with mock.patch.object(celery_executor.app, 'backend',
                                   mock_backend):
                mock_session = mock_backend.ResultSession.return_value  # pylint: disable=no-member
                mock_session.query.return_value.filter.return_value.all.return_value = [
                    mock.MagicMock(
                        **{
                            "to_dict.return_value": {
                                "status": "SUCCESS",
                                "task_id": "123"
                            }
                        })
                ]

        fetcher = BulkStateFetcher()
        result = fetcher.get_many([
            mock.MagicMock(task_id="123"),
            mock.MagicMock(task_id="456"),
        ])

        self.assertEqual(result, {
            '123': ('SUCCESS', None),
            '456': ("PENDING", None)
        })
    def test_should_support_db_backend(self, mock_session):
        with _prepare_app():
            mock_backend = DatabaseBackend(app=celery_executor.app,
                                           url="sqlite3://")

            with mock.patch.object(
                    celery_executor.app, 'backend',
                    mock_backend), self.assertLogs(
                        "airflow.executors.celery_executor.BulkStateFetcher",
                        level="DEBUG") as cm:
                mock_session = mock_backend.ResultSession.return_value  # pylint: disable=no-member
                mock_session.query.return_value.filter.return_value.all.return_value = [
                    mock.MagicMock(
                        **{
                            "to_dict.return_value": {
                                "status": "SUCCESS",
                                "task_id": "123"
                            }
                        })
                ]

                fetcher = BulkStateFetcher()
                result = fetcher.get_many([
                    mock.MagicMock(task_id="123"),
                    mock.MagicMock(task_id="456"),
                ])

        assert result == {'123': ('SUCCESS', None), '456': ("PENDING", None)}
        assert [
            'DEBUG:airflow.executors.celery_executor.BulkStateFetcher:Fetched 2 state(s) for 2 task(s)'
        ] == cm.output
    def test_should_support_kv_backend(self, mock_mget):
        with _prepare_app():
            mock_backend = BaseKeyValueStoreBackend(app=celery_executor.app)
            with mock.patch.object(
                    celery_executor.app, 'backend',
                    mock_backend), self.assertLogs(
                        "airflow.executors.celery_executor.BulkStateFetcher",
                        level="DEBUG") as cm:
                fetcher = BulkStateFetcher()
                result = fetcher.get_many([
                    mock.MagicMock(task_id="123"),
                    mock.MagicMock(task_id="456"),
                ])

        # Assert called - ignore order
        mget_args, _ = mock_mget.call_args
        assert set(mget_args[0]) == {
            b'celery-task-meta-456', b'celery-task-meta-123'
        }
        mock_mget.assert_called_once_with(mock.ANY)

        assert result == {'123': ('SUCCESS', None), '456': ("PENDING", None)}
        assert [
            'DEBUG:airflow.executors.celery_executor.BulkStateFetcher:Fetched 2 state(s) for 2 task(s)'
        ] == cm.output
Beispiel #4
0
    def test_should_support_base_backend(self):
        with _prepare_app():
            mock_backend = mock.MagicMock(autospec=BaseBackend)

            with mock.patch.object(celery_executor.app, 'backend', mock_backend):
                fetcher = BulkStateFetcher(1)
                result = fetcher.get_many([
                    ClassWithCustomAttributes(task_id="123", state='SUCCESS'),
                    ClassWithCustomAttributes(task_id="456", state="PENDING"),
                ])

        self.assertEqual(result, {'123': 'SUCCESS', '456': "PENDING"})
Beispiel #5
0
    def test_should_support_kv_backend(self, mock_mget):
        with _prepare_app():
            mock_backend = BaseKeyValueStoreBackend(app=celery_executor.app)
            with mock.patch.object(celery_executor.app, 'backend', mock_backend):
                fetcher = BulkStateFetcher()
                result = fetcher.get_many([
                    mock.MagicMock(task_id="123"),
                    mock.MagicMock(task_id="456"),
                ])

        # Assert called - ignore order
        mget_args, _ = mock_mget.call_args
        self.assertEqual(set(mget_args[0]), {b'celery-task-meta-456', b'celery-task-meta-123'})
        mock_mget.assert_called_once_with(mock.ANY)

        self.assertEqual(result, {'123': 'SUCCESS', '456': "PENDING"})
    def test_should_support_base_backend(self):
        with _prepare_app():
            mock_backend = mock.MagicMock(autospec=BaseBackend)

            with mock.patch.object(
                    celery_executor.app, 'backend',
                    mock_backend), self.assertLogs(
                        "airflow.executors.celery_executor.BulkStateFetcher",
                        level="DEBUG") as cm:
                fetcher = BulkStateFetcher(1)
                result = fetcher.get_many([
                    ClassWithCustomAttributes(task_id="123", state='SUCCESS'),
                    ClassWithCustomAttributes(task_id="456", state="PENDING"),
                ])

        assert result == {'123': ('SUCCESS', None), '456': ("PENDING", None)}
        assert [
            'DEBUG:airflow.executors.celery_executor.BulkStateFetcher:Fetched 2 state(s) for 2 task(s)'
        ] == cm.output