Update transports to send and receive data instead of strings (#42)
* Update transports to send and receive data instead of strings * Rename MockTransport helper methods * Adopt Task.sleep(for:...) instead of variant taking nanoseconds * Update comment * Revert inadvertant change to visibility of send method
This commit is contained in:
Mattt
GitHub
parent
de829ee648
commit
a9bcbf9ad9
1 changed files with 138 additions and 147 deletions
|
|
@ -6,10 +6,28 @@ import Logging
|
|||
/// Mock transport for testing
|
||||
actor MockTransport: Transport {
|
||||
var logger: Logger
|
||||
|
||||
let encoder = JSONEncoder()
|
||||
let decoder = JSONDecoder()
|
||||
|
||||
var isConnected = false
|
||||
private(set) var sentMessages: [String] = []
|
||||
private var messagesToReceive: [String] = []
|
||||
private var messageStreamContinuation: AsyncThrowingStream<String, Swift.Error>.Continuation?
|
||||
|
||||
private(set) var sentData: [Data] = []
|
||||
var sentMessages: [String] {
|
||||
return sentData.compactMap { data in
|
||||
guard let string = String(data: data, encoding: .utf8) else {
|
||||
logger.error("Failed to decode sent data as UTF-8")
|
||||
return nil
|
||||
}
|
||||
return string
|
||||
}
|
||||
}
|
||||
|
||||
private var dataToReceive: [Data] = []
|
||||
private(set) var receivedMessages: [String] = []
|
||||
|
||||
private var dataStreamContinuation: AsyncThrowingStream<Data, Swift.Error>.Continuation?
|
||||
|
||||
var shouldFailConnect = false
|
||||
var shouldFailSend = false
|
||||
|
||||
|
|
@ -17,85 +35,39 @@ actor MockTransport: Transport {
|
|||
self.logger = logger
|
||||
}
|
||||
|
||||
func connect() async throws {
|
||||
public func connect() async throws {
|
||||
if shouldFailConnect {
|
||||
throw Error.transportError(POSIXError(.ECONNREFUSED))
|
||||
}
|
||||
isConnected = true
|
||||
}
|
||||
|
||||
func disconnect() async {
|
||||
public func disconnect() async {
|
||||
isConnected = false
|
||||
messageStreamContinuation?.finish()
|
||||
messageStreamContinuation = nil
|
||||
dataStreamContinuation?.finish()
|
||||
dataStreamContinuation = nil
|
||||
}
|
||||
|
||||
func send<T: Encodable & Sendable>(_ message: T) async throws {
|
||||
public func send(_ message: Data) async throws {
|
||||
if shouldFailSend {
|
||||
throw Error.transportError(POSIXError(.EIO))
|
||||
}
|
||||
let data = try JSONEncoder().encode(message)
|
||||
let str = String(data: data, encoding: .utf8)!
|
||||
sentMessages.append(str)
|
||||
sentData.append(message)
|
||||
}
|
||||
|
||||
func receive() -> AsyncThrowingStream<String, Swift.Error> {
|
||||
return AsyncThrowingStream<String, Swift.Error> { continuation in
|
||||
messageStreamContinuation = continuation
|
||||
// Send any queued messages
|
||||
for message in messagesToReceive {
|
||||
public func receive() -> AsyncThrowingStream<Data, Swift.Error> {
|
||||
return AsyncThrowingStream<Data, Swift.Error> { continuation in
|
||||
dataStreamContinuation = continuation
|
||||
for message in dataToReceive {
|
||||
continuation.yield(message)
|
||||
if let string = String(data: message, encoding: .utf8) {
|
||||
receivedMessages.append(string)
|
||||
}
|
||||
}
|
||||
messagesToReceive.removeAll()
|
||||
dataToReceive.removeAll()
|
||||
}
|
||||
}
|
||||
|
||||
func queueRequest<M: MCP.Method>(_ request: Request<M>) throws {
|
||||
let data = try JSONEncoder().encode(request)
|
||||
let str = String(data: data, encoding: .utf8)!
|
||||
if let continuation = messageStreamContinuation {
|
||||
continuation.yield(str)
|
||||
} else {
|
||||
sentMessages.append(str)
|
||||
}
|
||||
}
|
||||
|
||||
func queueResponse<M: MCP.Method>(_ response: Response<M>) throws {
|
||||
let data = try JSONEncoder().encode(response)
|
||||
let str = String(data: data, encoding: .utf8)!
|
||||
if let continuation = messageStreamContinuation {
|
||||
continuation.yield(str)
|
||||
} else {
|
||||
messagesToReceive.append(str)
|
||||
}
|
||||
}
|
||||
|
||||
func queueNotification<N: MCP.Notification>(_ notification: Message<N>) throws {
|
||||
let data = try JSONEncoder().encode(notification)
|
||||
let str = String(data: data, encoding: .utf8)!
|
||||
if let continuation = messageStreamContinuation {
|
||||
continuation.yield(str)
|
||||
} else {
|
||||
messagesToReceive.append(str)
|
||||
}
|
||||
}
|
||||
|
||||
func getLastSentMessage<T: Decodable>() -> T? {
|
||||
print("SENT:", sentMessages)
|
||||
guard let lastMessage = sentMessages.last else { return nil }
|
||||
do {
|
||||
let data = lastMessage.data(using: .utf8)!
|
||||
return try JSONDecoder().decode(T.self, from: data)
|
||||
} catch {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func clearMessages() {
|
||||
sentMessages.removeAll()
|
||||
messagesToReceive.removeAll()
|
||||
}
|
||||
|
||||
func setFailConnect(_ shouldFail: Bool) {
|
||||
shouldFailConnect = shouldFail
|
||||
}
|
||||
|
|
@ -103,4 +75,37 @@ actor MockTransport: Transport {
|
|||
func setFailSend(_ shouldFail: Bool) {
|
||||
shouldFailSend = shouldFail
|
||||
}
|
||||
|
||||
func queue<M: MCP.Method>(request: Request<M>) throws {
|
||||
let data = try encoder.encode(request)
|
||||
if let continuation = dataStreamContinuation {
|
||||
continuation.yield(data)
|
||||
} else {
|
||||
dataToReceive.append(data)
|
||||
}
|
||||
}
|
||||
|
||||
func queue<M: MCP.Method>(response: Response<M>) throws {
|
||||
let data = try encoder.encode(response)
|
||||
dataToReceive.append(data)
|
||||
}
|
||||
|
||||
func queue<N: MCP.Notification>(notification: Message<N>) throws {
|
||||
let data = try encoder.encode(notification)
|
||||
dataToReceive.append(data)
|
||||
}
|
||||
|
||||
func decodeLastSentMessage<T: Decodable>() -> T? {
|
||||
guard let lastMessage = sentData.last else { return nil }
|
||||
do {
|
||||
return try decoder.decode(T.self, from: lastMessage)
|
||||
} catch {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func clearMessages() {
|
||||
sentData.removeAll()
|
||||
dataToReceive.removeAll()
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue