001/*
002 * SPDX-License-Identifier: Apache-2.0
003 *
004 * Copyright 2025-2026 The Enola <https://enola.dev> Authors
005 *
006 * Licensed under the Apache License, Version 2.0 (the "License");
007 * you may not use this file except in compliance with the License.
008 * You may obtain a copy of the License at
009 *
010 *     https://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package dev.enola.ai.adk.test;
019
020import com.google.adk.agents.BaseAgent;
021import com.google.adk.agents.InvocationContext;
022import com.google.adk.events.Event;
023import com.google.genai.types.Content;
024import com.google.genai.types.Part;
025
026import io.reactivex.rxjava3.core.Flowable;
027
028import java.util.List;
029import java.util.function.Supplier;
030
031public class MockAgent extends BaseAgent {
032
033    // https://github.com/google/adk-java/blob/main/core/src/test/java/com/google/adk/testing/TestBaseAgent.java
034    //   is similar to this, but it's in src/test and thus not available...
035    //   TODO Propose a related refactoring!
036
037    private final Supplier<Flowable<Event>> eventSupplier;
038    private String reply;
039
040    public MockAgent(String reply) {
041        super("Mock", "Mock Agent for Testing", List.of(), null, null);
042        this.reply = reply;
043
044        this.eventSupplier =
045                () ->
046                        Flowable.just(
047                                Event.builder()
048                                        .id(Event.generateEventId())
049                                        .author("MockAgent")
050                                        .content(Content.fromParts(Part.fromText(this.reply)))
051                                        .build());
052    }
053
054    public void replyWith(String text) {
055        this.reply = text;
056    }
057
058    @Override
059    protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
060        return eventSupplier.get();
061    }
062
063    @Override
064    protected Flowable<Event> runLiveImpl(InvocationContext invocationContext) {
065        return eventSupplier.get();
066    }
067}