47 lines
1.6 KiB
Swift
47 lines
1.6 KiB
Swift
import XCTest
|
|
@testable import MLX_Server
|
|
|
|
@MainActor
|
|
final class ChatViewModelTests: XCTestCase {
|
|
func testGemmaChatViewModelSendProducesAssistantReply() async throws {
|
|
let modelManager = ModelManager()
|
|
let config = try XCTUnwrap(ModelConfig.resolve("gemma"))
|
|
await modelManager.loadModel(config)
|
|
defer { modelManager.unloadModel() }
|
|
|
|
XCTAssertTrue(modelManager.isReady)
|
|
|
|
let viewModel = ChatViewModel(modelManager: modelManager)
|
|
viewModel.inputText = "Say hello in one word."
|
|
viewModel.send()
|
|
|
|
XCTAssertTrue(viewModel.isGenerating)
|
|
|
|
try await waitUntil(timeoutSeconds: 15) {
|
|
!viewModel.isGenerating
|
|
}
|
|
|
|
XCTAssertEqual(viewModel.conversation.messages.count, 2)
|
|
XCTAssertEqual(viewModel.conversation.messages[0].role, .user)
|
|
XCTAssertEqual(viewModel.conversation.messages[0].content, "Say hello in one word.")
|
|
XCTAssertEqual(viewModel.conversation.messages[1].role, .assistant)
|
|
XCTAssertFalse(viewModel.conversation.messages[1].sessionContent.isEmpty)
|
|
XCTAssertGreaterThan(viewModel.promptTokens, 0)
|
|
}
|
|
|
|
private func waitUntil(
|
|
timeoutSeconds: TimeInterval,
|
|
intervalNanoseconds: UInt64 = 100_000_000,
|
|
condition: @escaping @MainActor () -> Bool
|
|
) async throws {
|
|
let deadline = Date().addingTimeInterval(timeoutSeconds)
|
|
while Date() < deadline {
|
|
if condition() {
|
|
return
|
|
}
|
|
try await Task.sleep(nanoseconds: intervalNanoseconds)
|
|
}
|
|
XCTFail("Condition not met before timeout")
|
|
}
|
|
}
|