import os

from langchain_core.runnables import RunnableLambda
from streamlit.testing.v1 import AppTest


os.environ["STREAM_DELAY_SECONDS"] = "0"


def stream_reply(question: str):
    for word in f"LangChain streamed this response for: {question}".split():
        yield word + " "


chain = RunnableLambda(stream_reply)
chunks = list(chain.stream("show token flow"))
assert len(chunks) > 4, chunks
assert "".join(chunks).strip() == "LangChain streamed this response for: show token flow"
print(f"chunk_count={len(chunks)}")
print("stream_text=" + "".join(chunks).strip())

app = AppTest.from_file("app.py")
app.run(timeout=10)
assert not app.exception, app.exception
app.chat_input[0].set_value("show token flow").run(timeout=10)
assert not app.exception, app.exception

messages = [
    block.markdown[0].value
    for block in app.chat_message
    if getattr(block, "markdown", None)
]
assert "show token flow" in messages, messages
assistant_reply = next(
    message
    for message in messages
    if "LangChain streamed this response for: show token flow" in message
)
print(f"chat_messages={len(messages)}")
print("assistant_reply=" + assistant_reply)