be/src/pipeline/exec/mock_operator.h

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#pragma once

#include <list>

#include "pipeline/exec/operator.h"
#include "vec/core/block.h"

namespace doris::pipeline {

#ifdef BE_TEST
class MockLocalState final : public doris::pipeline::PipelineXLocalState<FakeSharedState> {
public:
    using Base = doris::pipeline::PipelineXLocalState<FakeSharedState>;
    ENABLE_FACTORY_CREATOR(MockLocalState);

    MockLocalState(RuntimeState* state, OperatorXBase* parent) : Base(state, parent) {}
    ~MockLocalState() override = default;

private:
    friend class MockOperatorX;
};

class MockOperatorX final : public OperatorX<MockLocalState> {
public:
    ENABLE_FACTORY_CREATOR(MockOperatorX);
    MockOperatorX() = default;

    Status get_block(RuntimeState* state, vectorized::Block* block, bool* eos) override {
        if (_outout_blocks.empty()) {
            *eos = true;
            return Status::OK();
        }

        *eos = false;
        block->swap(_outout_blocks.front());
        _outout_blocks.pop_front();
        if (_outout_blocks.empty()) {
            *eos = true;
        }
        return Status::OK();
    }

    [[nodiscard]] bool is_source() const override { return true; }
    friend class MockLocalState;
    std::list<vectorized::Block> _outout_blocks;
};

#endif

} // namespace doris::pipeline