claude-code/src/utils/__tests__/stream.test.ts

import { describe, expect, test } from 'bun:test'
import { Stream } from '../stream'

describe('Stream', () => {
  test('enqueue then read resolves with the value', async () => {
    const stream = new Stream<number>()
    stream[Symbol.asyncIterator]()
    stream.enqueue(42)
    const result = await stream.next()
    expect(result).toEqual({ done: false, value: 42 })
  })

  test('enqueue multiple then drain in order', async () => {
    const stream = new Stream<string>()
    stream[Symbol.asyncIterator]()
    stream.enqueue('a')
    stream.enqueue('b')
    stream.enqueue('c')
    expect(await stream.next()).toEqual({ done: false, value: 'a' })
    expect(await stream.next()).toEqual({ done: false, value: 'b' })
    expect(await stream.next()).toEqual({ done: false, value: 'c' })
  })

  test('next() blocks until enqueue provides a value', async () => {
    const stream = new Stream<number>()
    stream[Symbol.asyncIterator]()
    const promise = stream.next()
    // Not resolved yet — enqueue after a microtask
    stream.enqueue(99)
    const result = await promise
    expect(result).toEqual({ done: false, value: 99 })
  })

  test('done() resolves pending reader with done:true', async () => {
    const stream = new Stream<number>()
    stream[Symbol.asyncIterator]()
    const promise = stream.next()
    stream.done()
    expect(await promise).toEqual({ done: true, value: undefined })
  })

  test('done() with no pending reader — subsequent next returns done:true', async () => {
    const stream = new Stream<number>()
    stream[Symbol.asyncIterator]()
    stream.done()
    expect(await stream.next()).toEqual({ done: true, value: undefined })
  })

  test('error() rejects pending reader', async () => {
    const stream = new Stream<number>()
    stream[Symbol.asyncIterator]()
    const promise = stream.next()
    stream.error(new Error('boom'))
    expect(promise).rejects.toThrow('boom')
  })

  test('error() after done — hasError is set but next returns done:true (isDone checked first)', async () => {
    const stream = new Stream<number>()
    stream[Symbol.asyncIterator]()
    stream.done()
    stream.error(new Error('late error'))
    // next() checks isDone before hasError, so it returns done:true
    expect(await stream.next()).toEqual({ done: true, value: undefined })
  })

  test('enqueue after done — queue is checked before isDone, value is consumed', async () => {
    const stream = new Stream<number>()
    stream[Symbol.asyncIterator]()
    stream.done()
    stream.enqueue(1)
    // next() checks queue.length > 0 first, so enqueued value is returned
    expect(await stream.next()).toEqual({ done: false, value: 1 })
    // After draining queue, done takes effect
    expect(await stream.next()).toEqual({ done: true, value: undefined })
  })

  test('return() marks stream as done and calls returned callback', async () => {
    let called = false
    const stream = new Stream<number>(() => {
      called = true
    })
    stream[Symbol.asyncIterator]()
    const result = await stream.return()
    expect(result).toEqual({ done: true, value: undefined })
    expect(called).toBe(true)
    // Subsequent next returns done
    expect(await stream.next()).toEqual({ done: true, value: undefined })
  })

  test('return() without callback still works', async () => {
    const stream = new Stream<number>()
    stream[Symbol.asyncIterator]()
    const result = await stream.return()
    expect(result).toEqual({ done: true, value: undefined })
  })

  test('Symbol.asyncIterator throws on second call', () => {
    const stream = new Stream<number>()
    stream[Symbol.asyncIterator]()
    expect(() => stream[Symbol.asyncIterator]()).toThrow(
      'Stream can only be iterated once',
    )
  })

  test('for-await-of iteration drains queued values then ends', async () => {
    const stream = new Stream<string>()
    stream.enqueue('x')
    stream.enqueue('y')
    stream.done()
    const results: string[] = []
    for await (const value of stream) {
      results.push(value)
    }
    expect(results).toEqual(['x', 'y'])
  })

  test('for-await-of blocks until done', async () => {
    const stream = new Stream<number>()
    const results: number[] = []

    const iterPromise = (async () => {
      for await (const v of stream) {
        results.push(v)
      }
    })()

    // Enqueue after a tick
    await Promise.resolve()
    stream.enqueue(1)
    stream.enqueue(2)
    stream.done()

    await iterPromise
    expect(results).toEqual([1, 2])
  })

  test('error during for-await-of rejects the loop', async () => {
    const stream = new Stream<number>()
    const iterPromise = (async () => {
      for await (const _ of stream) {
        // will error before any value
      }
    })()
    stream.error(new Error('stream broken'))
    expect(iterPromise).rejects.toThrow('stream broken')
  })

  test('concurrent enqueue from multiple sources does not lose data', async () => {
    const stream = new Stream<number>()
    // Rapid sequential enqueue
    for (let i = 0; i < 100; i++) {
      stream.enqueue(i)
    }
    stream.done()

    const results: number[] = []
    for await (const v of stream) {
      results.push(v)
    }
    expect(results.length).toBe(100)
    expect(results[0]).toBe(0)
    expect(results[99]).toBe(99)
  })
})