327 lines
8.8 KiB
Text
327 lines
8.8 KiB
Text
import Std.Async
|
||
import Std.Sync
|
||
|
||
open Std.Async
|
||
|
||
-- Test basic message reception from multiple channels
|
||
def testSimpleMessages : Async Unit := do
|
||
let channelA ← Std.Channel.new (α := Nat)
|
||
let channelB ← Std.Channel.new (α := Nat)
|
||
let channelC ← Std.Channel.new (α := Nat)
|
||
|
||
let channel := Std.StreamMap.ofArray #[
|
||
("a", channelA),
|
||
("b", channelB),
|
||
("c", channelC),
|
||
]
|
||
|
||
await (← channelC.send 1)
|
||
let (name, message) ← channel.recv
|
||
assert! name == "c" && message == 1
|
||
|
||
await (← channelA.send 2)
|
||
let (name, message) ← channel.recv
|
||
assert! name == "a" && message == 2
|
||
|
||
await (← channelB.send 3)
|
||
let (name, message) ← channel.recv
|
||
assert! name == "b" && message == 3
|
||
|
||
#eval testSimpleMessages.block
|
||
|
||
-- Test empty StreamMap
|
||
def testEmpty : Async Unit := do
|
||
let stream : Std.StreamMap String Nat := Std.StreamMap.empty
|
||
|
||
assert! stream.isEmpty
|
||
assert! stream.size == 0
|
||
assert! stream.keys.size == 0
|
||
|
||
#eval testEmpty.block
|
||
|
||
-- Test register and unregister operations
|
||
def testRegisterUnregister : Async Unit := do
|
||
let channelA ← Std.Channel.new (α := Nat)
|
||
let channelB ← Std.Channel.new (α := Nat)
|
||
|
||
let stream := Std.StreamMap.empty.register "a" channelA
|
||
|
||
assert! stream.contains "a"
|
||
assert! not (stream.contains "b")
|
||
assert! stream.size == 1
|
||
|
||
let stream := stream.register "b" channelB
|
||
assert! stream.contains "a" && stream.contains "b"
|
||
assert! stream.size == 2
|
||
|
||
let stream := stream.unregister "a"
|
||
assert! not (stream.contains "a")
|
||
assert! stream.contains "b"
|
||
assert! stream.size == 1
|
||
|
||
let stream := stream.unregister "b"
|
||
assert! stream.isEmpty
|
||
|
||
#eval testRegisterUnregister.block
|
||
|
||
-- Test replacing existing stream with same name
|
||
def testRegisterReplace : Async Unit := do
|
||
let channelA1 ← Std.Channel.new (α := Nat)
|
||
let channelA2 ← Std.Channel.new (α := Nat)
|
||
|
||
let stream := Std.StreamMap.empty.register "a" channelA1
|
||
assert! stream.size == 1
|
||
|
||
-- Register with same name should replace
|
||
let stream := stream.register "a" channelA2
|
||
assert! stream.size == 1
|
||
assert! stream.contains "a"
|
||
|
||
-- Send to new channel
|
||
await (← channelA2.send 42)
|
||
let (name, message) ← stream.recv
|
||
assert! name == "a" && message == 42
|
||
|
||
#eval testRegisterReplace.block
|
||
|
||
-- Test keys functionality
|
||
def testKeys : Async Unit := do
|
||
let channelA ← Std.Channel.new (α := Nat)
|
||
let channelB ← Std.Channel.new (α := Nat)
|
||
let channelC ← Std.Channel.new (α := Nat)
|
||
|
||
let stream := Std.StreamMap.ofArray #[
|
||
("c", .mk channelC),
|
||
("a", .mk channelA),
|
||
("b", .mk channelB),
|
||
]
|
||
|
||
let keys := stream.keys
|
||
assert! keys.size == 3
|
||
assert! keys.contains "a"
|
||
assert! keys.contains "b"
|
||
assert! keys.contains "c"
|
||
|
||
#eval testKeys.block
|
||
|
||
-- Test get? functionality
|
||
def testGet : Async Unit := do
|
||
let channelA ← Std.Channel.new (α := Nat)
|
||
let channelB ← Std.Channel.new (α := Nat)
|
||
|
||
let stream := Std.StreamMap.ofArray #[
|
||
("a", .mk channelA),
|
||
("b", .mk channelB),
|
||
]
|
||
|
||
let selectorA := stream.get? "a"
|
||
let selectorC := stream.get? "c"
|
||
|
||
assert! selectorA.isSome
|
||
assert! selectorC.isNone
|
||
|
||
#eval testGet.block
|
||
|
||
-- Test filterByName functionality
|
||
def testFilterByName : Async Unit := do
|
||
let channelA ← Std.Channel.new (α := Nat)
|
||
let channelB ← Std.Channel.new (α := Nat)
|
||
let channelC ← Std.Channel.new (α := Nat)
|
||
|
||
let stream := Std.StreamMap.ofArray #[
|
||
("prefix_a", .mk channelA),
|
||
("prefix_b", .mk channelB),
|
||
("other_c", .mk channelC),
|
||
]
|
||
|
||
let filtered := stream.filterByName (fun name => name.startsWith "prefix_")
|
||
|
||
assert! filtered.size == 2
|
||
assert! filtered.contains "prefix_a"
|
||
assert! filtered.contains "prefix_b"
|
||
assert! not (filtered.contains "other_c")
|
||
|
||
#eval testFilterByName.block
|
||
|
||
-- Test toArray functionality
|
||
def testToArray : Async Unit := do
|
||
let channelA ← Std.Channel.new (α := Nat)
|
||
let channelB ← Std.Channel.new (α := Nat)
|
||
|
||
let stream := Std.StreamMap.ofArray #[
|
||
("a", .mk channelA),
|
||
("b", .mk channelB),
|
||
]
|
||
|
||
let array := stream.toArray
|
||
assert! array.size == 2
|
||
|
||
let names := array.map (·.1)
|
||
assert! names.contains "a"
|
||
assert! names.contains "b"
|
||
|
||
#eval testToArray.block
|
||
|
||
-- Test multiple messages from same channel
|
||
def testMultipleFromSame : Async Unit := do
|
||
let channelA ← Std.Channel.new (α := Nat)
|
||
let channelB ← Std.Channel.new (α := Nat)
|
||
|
||
let stream := Std.StreamMap.ofArray #[
|
||
("a", .mk channelA),
|
||
("b", .mk channelB),
|
||
]
|
||
|
||
-- Send multiple messages to same channel
|
||
await (← channelA.send 1)
|
||
await (← channelA.send 2)
|
||
await (← channelB.send 10)
|
||
|
||
let (name1, msg1) ← stream.recv
|
||
let (name2, msg2) ← stream.recv
|
||
let (name3, msg3) ← stream.recv
|
||
|
||
-- Should receive all messages, order may vary but names should match sources
|
||
assert! (name1 == "a" && msg1 == 1) || (name1 == "a" && msg1 == 2) || (name1 == "b" && msg1 == 10)
|
||
assert! (name2 == "a" && msg2 == 1) || (name2 == "a" && msg2 == 2) || (name2 == "b" && msg2 == 10)
|
||
assert! (name3 == "a" && msg3 == 1) || (name3 == "a" && msg3 == 2) || (name3 == "b" && msg3 == 10)
|
||
|
||
#eval testMultipleFromSame.block
|
||
|
||
-- Test interleaved messages from different channels
|
||
def testInterleavedMessages : Async Unit := do
|
||
let channelA ← Std.Channel.new (α := String)
|
||
let channelB ← Std.Channel.new (α := String)
|
||
let channelC ← Std.Channel.new (α := String)
|
||
|
||
let stream := Std.StreamMap.ofArray #[
|
||
("first", .mk channelA),
|
||
("second", .mk channelB),
|
||
("third", .mk channelC),
|
||
]
|
||
|
||
-- Send messages in specific order
|
||
await (← channelB.send "msg1")
|
||
await (← channelA.send "msg2")
|
||
await (← channelC.send "msg3")
|
||
await (← channelA.send "msg4")
|
||
|
||
let results ← (List.range 4).mapM (fun _ => stream.recv)
|
||
|
||
-- Verify we got all messages (order may vary)
|
||
let messages := results.map (·.2)
|
||
assert! messages.contains "msg1"
|
||
assert! messages.contains "msg2"
|
||
assert! messages.contains "msg3"
|
||
assert! messages.contains "msg4"
|
||
|
||
#eval testInterleavedMessages.block
|
||
|
||
-- Test with different data typez
|
||
def testDifferentTypes : Async Unit := do
|
||
let channelStr ← Std.Channel.new (α := String)
|
||
let channelBool ← Std.Channel.new (α := String)
|
||
|
||
let stream := Std.StreamMap.ofArray #[
|
||
("strings", .mk channelStr),
|
||
("bools", .mk channelBool),
|
||
]
|
||
|
||
await (← channelStr.send "hello")
|
||
await (← channelBool.send "world")
|
||
|
||
let (name1, msg1) ← stream.recv
|
||
let (name2, msg2) ← stream.recv
|
||
|
||
assert! ((name1 == "strings" && msg1 == "hello") || (name1 == "bools" && msg1 == "world"))
|
||
assert! ((name2 == "strings" && msg2 == "hello") || (name2 == "bools" && msg2 == "world"))
|
||
assert! name1 != name2
|
||
|
||
#eval testDifferentTypes.block
|
||
|
||
-- Test unregister during operation
|
||
def testUnregisterDuringOperation : Async Unit := do
|
||
let channelA ← Std.Channel.new (α := Nat)
|
||
let channelB ← Std.Channel.new (α := Nat)
|
||
let channelC ← Std.Channel.new (α := Nat)
|
||
|
||
let stream := Std.StreamMap.ofArray #[
|
||
("a", .mk channelA),
|
||
("b", .mk channelB),
|
||
("c", .mk channelC),
|
||
]
|
||
|
||
await (← channelA.send 1)
|
||
await (← channelB.send 2)
|
||
await (← channelC.send 3)
|
||
|
||
assert! (← stream.tryRecv).isSome
|
||
assert! (← stream.tryRecv).isSome
|
||
assert! (← stream.tryRecv).isSome
|
||
|
||
let stream := stream.unregister "b"
|
||
assert! not (stream.contains "b")
|
||
assert! stream.size == 2
|
||
|
||
let newChannelD ← Std.Channel.new (α := Nat)
|
||
let stream := stream.register "d" newChannelD
|
||
|
||
await (← newChannelD.send 4)
|
||
let (name2, msg2) ← stream.recv
|
||
|
||
assert! name2 == "d" && msg2 == 4
|
||
|
||
#eval testUnregisterDuringOperation.block
|
||
|
||
-- Test selector functionality
|
||
def testSelector : Async Unit := do
|
||
let channelA ← Std.Channel.new (α := Nat)
|
||
let channelB ← Std.Channel.new (α := Nat)
|
||
|
||
let stream := Std.StreamMap.ofArray #[
|
||
("a", .mk channelA),
|
||
("b", .mk channelB),
|
||
]
|
||
|
||
let selector ← stream.selector
|
||
|
||
await (← channelA.send 42)
|
||
|
||
let result ← Selectable.one #[.case selector (fun x => pure x)]
|
||
assert! result.1 == "a" && result.2 == 42
|
||
|
||
#eval testSelector.block
|
||
|
||
def testClose : Async Unit := do
|
||
let channelA ← Std.Channel.new (α := Nat)
|
||
let channelB ← Std.Channel.new (α := Nat)
|
||
|
||
let stream := Std.StreamMap.ofArray #[
|
||
("a", .mk channelA),
|
||
("b", .mk channelB),
|
||
]
|
||
|
||
stream.close
|
||
|
||
#eval testClose.block
|
||
|
||
-- Test large number of channels
|
||
def testManyChannels : Async Unit := do
|
||
let channels : Vector _ 128 ← Vector.ofFnM (fun _ => Std.Channel.new (α := Nat))
|
||
|
||
let streamArray := channels.mapIdx (fun i ch => (s!"channel_{i}", .mk ch))
|
||
let stream := Std.StreamMap.ofArray streamArray.toArray
|
||
|
||
assert! stream.size == 128
|
||
|
||
await (← channels[3].send 100)
|
||
await (← channels[7].send 200)
|
||
|
||
let (name1, msg1) ← stream.recv
|
||
let (name2, msg2) ← stream.recv
|
||
|
||
assert! ((name1 == "channel_3" && msg1 == 100) || (name1 == "channel_7" && msg1 == 200))
|
||
assert! ((name2 == "channel_3" && msg2 == 100) || (name2 == "channel_7" && msg2 == 200))
|
||
assert! name1 != name2
|
||
|
||
#eval testManyChannels.block
|