Sybil-2/ios/Tests/SybilPresentationTests/AgentInlineThinkingDotsTests.swift

import SwiftUI
import UIKit
import XCTest
@testable import Sybil

@MainActor
final class AgentInlineThinkingDotsTests: XCTestCase {
    func testPendingDotsReplaceWithTextAndHideOnDoneFailureOrCancel() async throws {
        let messageID = "temp-assistant-inline-dots-fixture"
        let date = Date(timeIntervalSince1970: 1_750_000_000)
        let user = Message(id: "temp-user-inline-dots-fixture", createdAt: date, role: .user, content: "@Sybil, summarize this.", name: nil)
        var reply = Message(
            id: messageID,
            createdAt: date,
            role: .assistant,
            content: "",
            name: nil,
            metadata: .object(["agentAside": .bool(true)])
        )
        var state = StreamingAssistantState()
        var presentation = Self.presentation(messageID: messageID, state: state)

        XCTAssertTrue(shouldShowInlineAgentPendingDots(message: reply, isSending: true, presentation: presentation))
        await capture(messages: [user, reply], presentation: presentation, sending: true, name: "inline-agent-pending-dancing-dots")

        XCTAssertEqual(state.receiveDelta(" \n"), .replace(" \n"))
        reply.content = state.segmentText
        presentation = Self.presentation(messageID: messageID, state: state)
        XCTAssertTrue(shouldShowInlineAgentPendingDots(message: reply, isSending: true, presentation: presentation))

        XCTAssertEqual(state.receiveDelta("Here is the Agent reply."), .replace(" \nHere is the Agent reply."))
        reply.content = state.segmentText
        presentation = Self.presentation(messageID: messageID, state: state)
        XCTAssertFalse(shouldShowInlineAgentPendingDots(message: reply, isSending: true, presentation: presentation))
        await capture(messages: [user, reply], presentation: presentation, sending: true, name: "inline-agent-first-visible-text-same-bubble")

        XCTAssertEqual(state.receiveDone("Completed Agent reply."), .replace("Completed Agent reply."))
        reply.content = "Completed Agent reply."
        presentation = Self.presentation(messageID: messageID, state: state)
        XCTAssertFalse(shouldShowInlineAgentPendingDots(message: reply, isSending: true, presentation: presentation))
        XCTAssertEqual(reply.id, messageID)

        var failed = StreamingAssistantState()
        let failurePresentation = Self.presentation(messageID: messageID, state: failed)
        XCTAssertTrue(shouldShowInlineAgentPendingDots(message: reply.withContent(""), isSending: true, presentation: failurePresentation))
        XCTAssertEqual(failed.receiveError(), .unchanged)
        XCTAssertFalse(shouldShowInlineAgentPendingDots(
            message: reply.withContent(""), isSending: true, presentation: Self.presentation(messageID: messageID, state: failed)
        ))

        var cancelled = StreamingAssistantState()
        let cancelPresentation = Self.presentation(messageID: messageID, state: cancelled)
        XCTAssertTrue(shouldShowInlineAgentPendingDots(message: reply.withContent(""), isSending: true, presentation: cancelPresentation))
        cancelled.receiveCancel()
        XCTAssertFalse(shouldShowInlineAgentPendingDots(
            message: reply.withContent(""), isSending: true, presentation: Self.presentation(messageID: messageID, state: cancelled)
        ))
    }

    private static func presentation(messageID: String, state: StreamingAssistantState) -> StreamingAssistantPresentation {
        StreamingAssistantPresentation(
            messageID: messageID,
            ephemeralTrace: state.ephemeralTrace,
            traceRevision: state.traceRevision,
            isAwaitingVisibleText: state.isAwaitingVisibleText
        )
    }

    private func capture(
        messages: [Message],
        presentation: StreamingAssistantPresentation,
        sending: Bool,
        name: String
    ) async {
        let host = UIHostingController(rootView: InlineAgentReplySnapshot(
            messages: messages,
            presentation: presentation,
            sending: sending
        ))
        let window = UIWindow(frame: CGRect(x: 0, y: 0, width: 390, height: 170))
        window.rootViewController = host
        window.makeKeyAndVisible()
        host.view.frame = window.bounds
        host.view.setNeedsLayout()
        host.view.layoutIfNeeded()
        try? await Task.sleep(for: .milliseconds(80))
        let image = UIGraphicsImageRenderer(bounds: host.view.bounds).image { _ in
            XCTAssertTrue(host.view.drawHierarchy(in: host.view.bounds, afterScreenUpdates: true))
        }
        if let data = image.pngData() {
            let attachment = XCTAttachment(data: data, uniformTypeIdentifier: "public.png")
            attachment.name = name
            attachment.lifetime = .keepAlways
            add(attachment)
        } else {
            XCTFail("Could not encode (name) screenshot")
        }
        window.isHidden = true
        window.rootViewController = nil
    }
}

private struct InlineAgentReplySnapshot: View {
    var messages: [Message]
    var presentation: StreamingAssistantPresentation
    var sending: Bool

    var body: some View {
        ZStack(alignment: .topLeading) {
            SybilTheme.background.ignoresSafeArea()
            VStack(alignment: .leading, spacing: 8) {
                Text("Inline @Sybil reply")
                    .font(.sybil(.caption))
                    .foregroundStyle(SybilTheme.textMuted)
                ForEach(messages) { message in
                    if message.isAgentAside {
                        AgentMessageRow(
                            message: message,
                            isSending: sending,
                            streamingAssistantPresentation: presentation
                        )
                    } else {
                        Text(message.content).foregroundStyle(SybilTheme.text)
                    }
                }
                Spacer(minLength: 0)
            }
            .padding(16)
        }
        .frame(width: 390, height: 170)
        .environment(\.colorScheme, .dark)
    }
}

private extension Message {
    func withContent(_ content: String) -> Message {
        var copy = self
        copy.content = content
        return copy
    }
}