Skip to content

Commit 212035c

Browse files
pblazejpblazej
committed
Extract tests
1 parent 1ea3f56 commit 212035c

File tree

1 file changed

+118
-129
lines changed

1 file changed

+118
-129
lines changed

‎Tests/LiveKitTests/Agent/TranscriptionTests.swift‎

Lines changed: 118 additions & 129 deletions
Original file line numberDiff line numberDiff line change
@@ -37,160 +37,149 @@ actor MessageCollector {
3737
}
3838

3939
class TranscriptionTests: LKTestCase, @unchecked Sendable {
40+
private var rooms: [Room] = []
41+
private var receiver: TranscriptionStreamReceiver!
42+
private var senderRoom: Room!
43+
private var messageCollector: MessageCollector!
44+
private var collectionTask: Task<Void, Never>!
45+
private var messageExpectation: XCTestExpectation!
46+
4047
// Same segment, same stream
4148
func testUpdates() async throws {
42-
let messageExpectation = expectation(description: "Receives all message updates")
43-
messageExpectation.expectedFulfillmentCount = 3
44-
4549
let segmentID = "test-segment"
46-
let topic = "lk.transcription"
47-
50+
let streamID = UUID().uuidString
4851
let testChunks = ["Hey", " there!", " What's up?"]
52+
let expectedContent = ["Hey", "Hey there!", "Hey there! What's up?"]
53+
54+
try await runTranscriptionTest(
55+
chunks: testChunks,
56+
segmentID: segmentID,
57+
streamID: streamID,
58+
expectedContent: expectedContent
59+
)
60+
}
4961

50-
try await withRooms([
51-
RoomTestingOptions(canSubscribe: true),
52-
RoomTestingOptions(canPublishData: true),
53-
]) { rooms in
54-
let receiverRoom = rooms[0]
55-
let senderRoom = rooms[1]
62+
// Same segment, different stream
63+
func testReplace() async throws {
64+
let segmentID = "test-segment"
65+
let testChunks = ["Hey", "Hey there!", "Hey there! What's up?"]
66+
let expectedContent = ["Hey", "Hey there!", "Hey there! What's up?"]
67+
68+
try await runTranscriptionTest(
69+
chunks: testChunks,
70+
segmentID: segmentID,
71+
streamID: nil,
72+
expectedContent: expectedContent
73+
)
74+
}
5675

57-
let receiver = TranscriptionStreamReceiver(room: receiverRoom)
58-
let messageStream = try await receiver.messages()
59-
let streamID = UUID().uuidString
76+
private func setupTestEnvironment(expectedCount: Int) async throws {
77+
messageExpectation = expectation(description: "Receives all message updates")
78+
messageExpectation.expectedFulfillmentCount = expectedCount
6079

61-
let messageCollector = MessageCollector()
80+
receiver = TranscriptionStreamReceiver(room: rooms[0])
81+
let messageStream = try await receiver.messages()
82+
messageCollector = MessageCollector()
83+
senderRoom = rooms[1]
6284

63-
let collectionTask = Task { @Sendable in
64-
var iterator = messageStream.makeAsyncIterator()
65-
while let message = await iterator.next() {
66-
await messageCollector.add(message)
67-
messageExpectation.fulfill()
68-
}
85+
collectionTask = Task { @Sendable in
86+
var iterator = messageStream.makeAsyncIterator()
87+
while let message = await iterator.next() {
88+
await self.messageCollector.add(message)
89+
self.messageExpectation.fulfill()
6990
}
91+
}
92+
}
7093

71-
for (index, chunk) in testChunks.enumerated() {
72-
let isLast = index == testChunks.count - 1
73-
74-
var attributes: [String: String] = [
75-
"lk.segment_id": segmentID,
76-
"lk.transcription_final": "false",
77-
]
94+
private func sendTranscriptionChunks(
95+
chunks: [String],
96+
segmentID: String,
97+
streamID: String? = nil,
98+
to room: Room
99+
) async throws {
100+
let topic = "lk.transcription"
78101

79-
if isLast {
80-
attributes["lk.transcription_final"] = "true"
81-
}
102+
for (index, chunk) in chunks.enumerated() {
103+
let isLast = index == chunks.count - 1
82104

83-
let options = StreamTextOptions(
84-
topic: topic,
85-
attributes: attributes,
86-
id: streamID
87-
)
105+
var attributes: [String: String] = [
106+
"lk.segment_id": segmentID,
107+
"lk.transcription_final": "false",
108+
]
88109

89-
try await senderRoom.localParticipant.sendText(chunk, options: options)
90-
try await Task.sleep(nanoseconds: 10_000_000)
110+
if isLast {
111+
attributes["lk.transcription_final"] = "true"
91112
}
92113

93-
await self.fulfillment(of: [messageExpectation], timeout: 5)
94-
collectionTask.cancel()
95-
96-
let updates = await messageCollector.getUpdates()
97-
XCTAssertEqual(updates.count, 3)
98-
XCTAssertEqual(updates[0].content, .agentTranscript("Hey"))
99-
XCTAssertEqual(updates[1].content, .agentTranscript("Hey there!"))
100-
XCTAssertEqual(updates[2].content, .agentTranscript("Hey there! What's up?"))
101-
102-
XCTAssertEqual(updates[0].id, segmentID)
103-
XCTAssertEqual(updates[1].id, segmentID)
104-
XCTAssertEqual(updates[2].id, segmentID)
105-
106-
let firstTimestamp = updates[0].timestamp
107-
XCTAssertEqual(updates[1].timestamp, firstTimestamp)
108-
XCTAssertEqual(updates[2].timestamp, firstTimestamp)
109-
110-
let messages = await messageCollector.getMessages()
111-
XCTAssertEqual(messages.count, 1)
112-
XCTAssertEqual(messages.keys[0], segmentID)
113-
XCTAssertEqual(messages.values[0].content, .agentTranscript("Hey there! What's up?"))
114-
XCTAssertEqual(messages.values[0].id, segmentID)
115-
XCTAssertEqual(messages.values[0].timestamp, firstTimestamp)
114+
let options = StreamTextOptions(
115+
topic: topic,
116+
attributes: attributes,
117+
id: streamID ?? UUID().uuidString
118+
)
119+
120+
try await room.localParticipant.sendText(chunk, options: options)
121+
try await Task.sleep(nanoseconds: 10_000_000)
116122
}
117123
}
118124

119-
// Same segment, different stream
120-
func testReplace() async throws {
121-
let messageExpectation = expectation(description: "Receives all message updates")
122-
messageExpectation.expectedFulfillmentCount = 3
125+
private func validateTranscriptionResults(
126+
updates: [ReceivedMessage],
127+
messages: OrderedDictionary<ReceivedMessage.ID, ReceivedMessage>,
128+
segmentID: String,
129+
expectedContent: [String]
130+
) {
131+
// Validate updates
132+
XCTAssertEqual(updates.count, expectedContent.count)
133+
for (index, expected) in expectedContent.enumerated() {
134+
XCTAssertEqual(updates[index].content, .agentTranscript(expected))
135+
XCTAssertEqual(updates[index].id, segmentID)
136+
}
123137

124-
let segmentID = "test-segment"
125-
let topic = "lk.transcription"
138+
// Validate timestamps are consistent
139+
let firstTimestamp = updates[0].timestamp
140+
for update in updates {
141+
XCTAssertEqual(update.timestamp, firstTimestamp)
142+
}
126143

127-
let testChunks = ["Hey", "Hey there!", "Hey there! What's up?"]
144+
// Validate final message
145+
XCTAssertEqual(messages.count, 1)
146+
XCTAssertEqual(messages.keys[0], segmentID)
147+
XCTAssertEqual(messages.values[0].content, .agentTranscript(expectedContent.last!))
148+
XCTAssertEqual(messages.values[0].id, segmentID)
149+
XCTAssertEqual(messages.values[0].timestamp, firstTimestamp)
150+
}
128151

152+
private func runTranscriptionTest(
153+
chunks: [String],
154+
segmentID: String,
155+
streamID: String? = nil,
156+
expectedContent: [String]
157+
) async throws {
129158
try await withRooms([
130159
RoomTestingOptions(canSubscribe: true),
131160
RoomTestingOptions(canPublishData: true),
132161
]) { rooms in
133-
let receiverRoom = rooms[0]
134-
let senderRoom = rooms[1]
135-
136-
let receiver = TranscriptionStreamReceiver(room: receiverRoom)
137-
let messageStream = try await receiver.messages()
138-
139-
let messageCollector = MessageCollector()
140-
141-
let collectionTask = Task { @Sendable in
142-
var iterator = messageStream.makeAsyncIterator()
143-
while let message = await iterator.next() {
144-
await messageCollector.add(message)
145-
messageExpectation.fulfill()
146-
}
147-
}
148-
149-
for (index, chunk) in testChunks.enumerated() {
150-
let isLast = index == testChunks.count - 1
151-
152-
var attributes: [String: String] = [
153-
"lk.segment_id": segmentID,
154-
"lk.transcription_final": "false",
155-
]
156-
157-
if isLast {
158-
attributes["lk.transcription_final"] = "true"
159-
}
160-
161-
let options = StreamTextOptions(
162-
topic: topic,
163-
attributes: attributes,
164-
id: UUID().uuidString
165-
)
166-
167-
try await senderRoom.localParticipant.sendText(chunk, options: options)
168-
try await Task.sleep(nanoseconds: 10_000_000)
169-
}
170-
171-
await self.fulfillment(of: [messageExpectation], timeout: 5)
172-
collectionTask.cancel()
173-
174-
let updates = await messageCollector.getUpdates()
175-
XCTAssertEqual(updates.count, 3)
176-
XCTAssertEqual(updates[0].content, .agentTranscript("Hey"))
177-
XCTAssertEqual(updates[1].content, .agentTranscript("Hey there!"))
178-
XCTAssertEqual(updates[2].content, .agentTranscript("Hey there! What's up?"))
179-
180-
XCTAssertEqual(updates[0].id, segmentID)
181-
XCTAssertEqual(updates[1].id, segmentID)
182-
XCTAssertEqual(updates[2].id, segmentID)
183-
184-
let firstTimestamp = updates[0].timestamp
185-
XCTAssertEqual(updates[1].timestamp, firstTimestamp)
186-
XCTAssertEqual(updates[2].timestamp, firstTimestamp)
187-
188-
let messages = await messageCollector.getMessages()
189-
XCTAssertEqual(messages.count, 1)
190-
XCTAssertEqual(messages.keys[0], segmentID)
191-
XCTAssertEqual(messages.values[0].content, .agentTranscript("Hey there! What's up?"))
192-
XCTAssertEqual(messages.values[0].id, segmentID)
193-
XCTAssertEqual(messages.values[0].timestamp, firstTimestamp)
162+
self.rooms = rooms
163+
try await self.setupTestEnvironment(expectedCount: expectedContent.count)
164+
try await self.sendTranscriptionChunks(
165+
chunks: chunks,
166+
segmentID: segmentID,
167+
streamID: streamID,
168+
to: self.senderRoom
169+
)
170+
171+
await self.fulfillment(of: [self.messageExpectation], timeout: 5)
172+
self.collectionTask.cancel()
173+
174+
let updates = await self.messageCollector.getUpdates()
175+
let messages = await self.messageCollector.getMessages()
176+
177+
self.validateTranscriptionResults(
178+
updates: updates,
179+
messages: messages,
180+
segmentID: segmentID,
181+
expectedContent: expectedContent
182+
)
194183
}
195184
}
196185
}

0 commit comments

Comments
 (0)