stream#EventEmitter TypeScript Examples

function doMockArchiverAsThrower() {
  let eventEmitter: EventEmitter;
  mockedArchiver.mockImplementationOnce(() => {
    eventEmitter = new EventEmitter();
    const archiver: Archiver = eventEmitter as Archiver;
    archiver.pipe = jest.fn();
    archiver.directory = jest.fn();
    archiver.finalize = jest.fn();
    return eventEmitter as Archiver;
  });
  return (err: unknown) => eventEmitter.emit('error', err);
}

class AudioResourceEvent extends EventEmitter {
  on<T extends keyof EventKeys>(event:T, callback:(...args:EventKeys[T])=>any):this{
    super.on(event, callback);
    return this;
  }
  off<T extends keyof EventKeys>(event:T, callback:(...args:EventKeys[T])=>any) {
    super.off(event, callback);
    return this;
  }
  once<T extends keyof EventKeys>(event:T, callback:(...args:EventKeys[T])=>any) {
    super.once(event, callback);
    return this;
  }
  addListener<T extends keyof EventKeys>(event:T, callback:(...args:EventKeys[T])=>any) {
    super.addListener(event, callback);
    return this;
  }
  removeListener<T extends keyof EventKeys>(event:T, callback:(...args:EventKeys[T])=>any) {
    super.removeListener(event, callback);
    return this;
  }
  removeAllListeners<T extends keyof EventKeys>(event:T) {
    super.removeAllListeners(event);
    return this;
  }
  emit<T extends keyof EventKeys>(event:T, ...args:EventKeys[T]){
    return super.emit(event, args);
  }
}

export class addOn extends EventEmitter {  
  on<T extends keyof EventKeys>(event:T, callback:(...args:EventKeys[T])=>any):this{
    super.on(event, callback);
    return this;
  }
  off<T extends keyof EventKeys>(event:T, callback:(...args:EventKeys[T])=>any) {
    super.off(event, callback);
    return this;
  }
  once<T extends keyof EventKeys>(event:T, callback:(...args:EventKeys[T])=>any) {
    super.once(event, callback);
    return this;
  }
  addListener<T extends keyof EventKeys>(event:T, callback:(...args:EventKeys[T])=>any) {
    super.addListener(event, callback);
    return this;
  }
  removeListener<T extends keyof EventKeys>(event:T, callback:(...args:EventKeys[T])=>any) {
    super.removeListener(event, callback);
    return this;
  }
  removeAllListeners<T extends keyof EventKeys>(event:T) {
    super.removeAllListeners(event);
    return this;
  }
  emit<T extends keyof EventKeys>(event:T, ...args:EventKeys[T]){
    return super.emit(event, args);
  }

  constructor(){
    super({captureRejections: false});
    try{
      fs.readdirSync(path.join(__dirname, "../../addon/"), {withFileTypes: true})
      .filter(d => d.isFile())
      .map(d => require("../../addon/" + d.name.slice(0, -3)))
      .filter(d => typeof d === "function")
      .forEach(d => {
        try {
          d(this);
        }
        catch{}
      });
    }
    catch{
    }
  }
}

export class DtlsSocket extends EventEmitter {
    socket?: dgram.Socket;
    state: HandshakeState;

    epoch: number;
    currentEpoch: CurrentEpoch;
    nextEpoch: NextEpoch;
    handshakeResendTimeout = 200;

    queuedApplicationData: Buffer[];
    serverCertificates: x509Certificate[];

    constructor() {
        super();

        this.state = HandshakeState.Established;

        this.epoch = 0;

        this.currentEpoch = {
            nextOutgoingSequence: 1,
            nextExpectedSequence: 1,
            previousSequenceWindowBitmask: 0,
            recordProtection: new NullRecordProtection
        };

        this.nextEpoch = {
            epoch: 1,
            state: HandshakeState.Initializing,
            nextOutgoingSequence: 1,
            nextPacketResendTime: 0,
            selectedCipherSuite: CipherSuite.TLS_NULL_WITH_NULL_NULL,
            recordProtection: new NullRecordProtection,
            handshake: undefined,
            cookie: Buffer.alloc(0),
            verificationStream: [],
            serverPublicKey: undefined,
            clientRandom: Buffer.alloc(32),
            serverRandom: Buffer.alloc(32),
            masterSecret: Buffer.alloc(32),
            serverVerification: Buffer.alloc(0),
            certificateFragments: [],
            certificatePayload: Buffer.alloc(0)
        };

        this.queuedApplicationData = [];
        this.serverCertificates = [];
    }

    async _connect(port: number, ip: string) {
        return new Promise<void>((resolve, reject) => {
            if (this.socket) {
                this.socket.close();
            }

            this.socket = dgram.createSocket("udp4");

            this.socket.on("message", msg => {
                const reader = HazelReader.from(msg);
                this.handleRecv(reader);
            });

            try {
                this.socket.on("error", err => {
                    reject(err);
                });

                this.socket.on("connect", resolve);

                this.socket.connect(port, ip);
            } catch (e) {
                reject(e);
            }
        });
    }

    async connect(port: number, ip: string) {
        await this._connect(port, ip);
        this.resetConnectionState();
        crypto.randomFillSync(this.nextEpoch.clientRandom, 0);
        this.sendClientHello();
    }

    async close() {
        this.socket?.close();
        this.resetConnectionState();
    }

    resetConnectionState() {
        this.state = HandshakeState.Established;

        this.epoch = 0;

        this.currentEpoch = {
            nextOutgoingSequence: 1,
            nextExpectedSequence: 1,
            previousSequenceWindowBitmask: 0,
            recordProtection: new NullRecordProtection
        };

        this.nextEpoch = {
            epoch: 1,
            state: HandshakeState.Initializing,
            nextOutgoingSequence: 1,
            nextPacketResendTime: 0,
            selectedCipherSuite: CipherSuite.TLS_NULL_WITH_NULL_NULL,
            recordProtection: new NullRecordProtection,
            handshake: undefined,
            cookie: Buffer.alloc(0),
            verificationStream: [],
            serverPublicKey: undefined,
            clientRandom: Buffer.alloc(32),
            serverRandom: Buffer.alloc(32),
            masterSecret: Buffer.alloc(32),
            serverVerification: Buffer.alloc(0),
            certificateFragments: [],
            certificatePayload: Buffer.alloc(0)
        };

        this.queuedApplicationData = [];
        this.serverCertificates = [];
    }

    setValidServerCertificates(certificates: x509Certificate[]) {
        this.serverCertificates.push(...certificates);
    }

    resendPacketsIfNeeded() {
        if (this.nextEpoch.state !== HandshakeState.Established) {
            const now = Date.now();
            if (now >= this.nextEpoch.nextPacketResendTime) {
                switch (this.nextEpoch.state) {
                    case HandshakeState.ExpectingServerHello:
                    case HandshakeState.ExpectingCertificate:
                    case HandshakeState.ExpectingServerKeyExchange:
                    case HandshakeState.ExpectingServerHelloDone:
                        this.sendClientHello();
                        break;
                    case HandshakeState.ExpectingChangeCipherSpec:
                    case HandshakeState.ExpectingFinished:
                        this.sendClientKeyExchangeFlight(true);
                        break;
                    default:
                        break;
                }
            }
        }
    }

    flushQueuedApplicationData() {
        for (const queuedData of this.queuedApplicationData) {
            const outgoingRecord = new RecordHeader(
                ContentType.ApplicationData,
                this.epoch,
                this.currentEpoch.nextOutgoingSequence,
                this.currentEpoch.recordProtection!.getEncryptedSize(queuedData.byteLength)
            );
            this.currentEpoch.nextOutgoingSequence++;

            const packet = HazelWriter.alloc(RecordHeader.size + outgoingRecord.length);
            packet.write(outgoingRecord);

            const output = this.currentEpoch.recordProtection!.encryptClientPlaintext(
                queuedData,
                outgoingRecord
            );
            packet.bytes(output);
            this.socket?.send(packet.buffer);
        }
        this.queuedApplicationData = [];
    }

    private createWireData(bytes: Buffer) {
        if (this.nextEpoch.state !== HandshakeState.Established) {
            const copy = Buffer.from(bytes);
            this.queuedApplicationData.push(copy);
            return Buffer.alloc(0);
        }

        this.flushQueuedApplicationData();

        const outgoingRecord = new RecordHeader(
            ContentType.ApplicationData,
            this.epoch,
            this.currentEpoch.nextOutgoingSequence,
            this.currentEpoch.recordProtection!.getEncryptedSize(bytes.byteLength)
        );
        this.currentEpoch.nextOutgoingSequence++;

        const packet = HazelWriter.alloc(RecordHeader.size + outgoingRecord.length);
        packet.write(outgoingRecord);
        packet.bytes(bytes);

        const output = this.currentEpoch.recordProtection!.encryptClientPlaintext(
            packet.buffer.slice(RecordHeader.size, RecordHeader.size + bytes.byteLength),
            outgoingRecord
        );

        output.copy(packet.buffer, RecordHeader.size);
        return packet.buffer;
    }

    send(bytes: Buffer, cb?: (err: Error|null, bytesWritten: number) => void) {
        const wireData = this.createWireData(bytes);
        if (wireData.byteLength) {
            this.socket?.send(wireData, cb);
        }
    }

    handleRecv(reader: HazelReader) {
        while (reader.left) {
            const record = reader.read(RecordHeader);
            const payloadReader = reader.bytes(record.length);

            if (record.contentType === ContentType.ApplicationData && this.nextEpoch.state !== HandshakeState.Established) {
                continue;
            }

            if (record.epoch !== this.epoch) {
                continue;
            }

            const windowIdx = this.currentEpoch.nextExpectedSequence - record.sequenceNumber - 1;
            const windowMask = 1 << windowIdx;
            if (record.sequenceNumber < this.currentEpoch.nextExpectedSequence) {
                if (windowIdx >= 64) {
                    continue;
                }

                if ((this.currentEpoch.previousSequenceWindowBitmask & windowMask) !== 0) {
                    continue;
                }
            }

            const decryptedPayload = this.currentEpoch.recordProtection.decryptCiphertextFromServer(payloadReader.buffer, record);

            if (!decryptedPayload) {
                return;
            }

            if (record.sequenceNumber >= this.currentEpoch.nextExpectedSequence) {
                const windowShift = record.sequenceNumber + 1 - this.currentEpoch.nextExpectedSequence;
                this.currentEpoch.previousSequenceWindowBitmask <<= windowShift;
                this.currentEpoch.nextExpectedSequence = record.sequenceNumber + 1;
            } else {
                this.currentEpoch.previousSequenceWindowBitmask |= windowMask;
            }

            const decryptedReader = HazelReader.from(decryptedPayload);
            this.handlePayload(record, decryptedReader);
        }
    }

    private handlePayload(record: RecordHeader, reader: HazelReader) {
        switch (record.contentType) {
            case ContentType.ChangeCipherSpec:
                if (this.nextEpoch.state !== HandshakeState.ExpectingChangeCipherSpec) {
                    return;
                }
                this.epoch = this.nextEpoch.epoch;
                this.currentEpoch.recordProtection = this.nextEpoch.recordProtection;
                this.currentEpoch.nextOutgoingSequence = this.nextEpoch.nextOutgoingSequence;
                this.currentEpoch.nextExpectedSequence = 1;
                this.currentEpoch.previousSequenceWindowBitmask = 0;

                this.nextEpoch.state = HandshakeState.ExpectingFinished;
                this.nextEpoch.selectedCipherSuite = CipherSuite.TLS_NULL_WITH_NULL_NULL;
                this.nextEpoch.recordProtection = new NullRecordProtection;
                this.nextEpoch.handshake = undefined;
                this.nextEpoch.cookie = Buffer.alloc(0);
                this.nextEpoch.verificationStream = [];
                this.nextEpoch.serverPublicKey = undefined;
                this.nextEpoch.serverRandom = Buffer.alloc(32);
                this.nextEpoch.clientRandom = Buffer.alloc(32);
                this.nextEpoch.masterSecret = Buffer.alloc(48);
                this.emit("ready");
                break;
            case ContentType.Alert:
                break;
            case ContentType.Handshake:
                this.handleHandshake(record, reader);
                break;
            case ContentType.ApplicationData:
                this.emit("message", reader.buffer);
                break;
        }
    }

    private handleHandshake(record: RecordHeader, reader: HazelReader) {
        while (reader.left) {
            const startOfHandshake = reader.cursor;
            const handshake = reader.read(Handshake);
            let payloadReader = reader.bytes(handshake.length);

            if (handshake.messageType !== HandshakeType.Certificate && (handshake.fragmentOffset !== 0 || handshake.fragmentLength !== handshake.length)) {
                continue;
            }

            switch (handshake.messageType) {
                case HandshakeType.HelloVerifyRequest:
                    if (this.nextEpoch.state !== HandshakeState.ExpectingServerHello) {
                        continue;
                    }

                    const helloVerifyRequest = payloadReader.read(HelloVerifyRequest);
                    this.nextEpoch.cookie = Buffer.alloc(helloVerifyRequest.cookie.byteLength);
                    helloVerifyRequest.cookie.copy(this.nextEpoch.cookie);

                    crypto.randomFillSync(this.nextEpoch.clientRandom);
                    this.sendClientHello();
                    break;
                case HandshakeType.ServerHello:
                    if (this.nextEpoch.state !== HandshakeState.ExpectingServerHello) {
                        continue;
                    } else if (handshake.messageSequence !== 1) {
                        continue;
                    }

                    const serverHello = payloadReader.read(ServerHello);

                    let selectedCipherSuite = CipherSuite.TLS_NULL_WITH_NULL_NULL;
                    for (const cipherSuite of serverHello.cipherSuites) {
                        switch (cipherSuite) {
                            case CipherSuite.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256:
                                this.nextEpoch.handshake = new X25519EcdheRsaSha256;
                                selectedCipherSuite = cipherSuite;
                                break;
                        }

                        if (selectedCipherSuite !== CipherSuite.TLS_NULL_WITH_NULL_NULL)
                            break;
                    }

                    this.nextEpoch.selectedCipherSuite = selectedCipherSuite;
                    serverHello.random.copy(this.nextEpoch.serverRandom);
                    this.nextEpoch.state = HandshakeState.ExpectingCertificate;
                    this.nextEpoch.certificateFragments = [];
                    this.nextEpoch.certificatePayload = Buffer.alloc(0);

                    this.nextEpoch.verificationStream.push(reader.buffer.slice(startOfHandshake, startOfHandshake + 12 + handshake.length));
                    break;
                case HandshakeType.Certificate:
                    if (this.nextEpoch.state !== HandshakeState.ExpectingCertificate) {
                        continue;
                    } else if (handshake.messageSequence !== 2) {
                        continue;
                    }

                    if (handshake.fragmentLength !== handshake.length) {
                        if (this.nextEpoch.certificatePayload.length !== handshake.length) {
                            this.nextEpoch.certificatePayload = Buffer.alloc(handshake.length);
                            this.nextEpoch.certificateFragments = [];
                        }
                        payloadReader.buffer.copy(this.nextEpoch.certificatePayload, handshake.fragmentOffset, 0, handshake.fragmentLength);
                        this.nextEpoch.certificateFragments.push({
                            offset: handshake.fragmentOffset,
                            length: handshake.fragmentLength
                        });
                        this.nextEpoch.certificateFragments.sort((a, b) => {
                            return a.offset - b.offset;
                        });

                        let currentOffset = 0;
                        let valid = false;
                        for (const fragmentRange of this.nextEpoch.certificateFragments) {
                            if (fragmentRange.offset !== currentOffset) {
                                valid = false;
                                break;
                            }
                            currentOffset += fragmentRange.length;
                        }
                        if (currentOffset !== this.nextEpoch.certificatePayload.length) {
                            valid = false;
                        }
                        if (!valid) {
                            continue;
                        }
                        this.nextEpoch.certificateFragments = [];
                        payloadReader = HazelReader.from(this.nextEpoch.certificatePayload);
                    }

                    const certificate = payloadReader.read(Certificate).certificates[0];

                    // todo: check if this.serverCertificates has this certificate.

                    if (certificate.publicKey.algo !== "rsaEncryption")
                        continue;

                    const fullCertificateHandshake = new Handshake(
                        handshake.messageType,
                        handshake.length,
                        handshake.messageSequence,
                        0,
                        handshake.length
                    );

                    const serializedCertificateHandshake = HazelWriter.alloc(12);
                    serializedCertificateHandshake.write(fullCertificateHandshake);
                    this.nextEpoch.verificationStream.push(serializedCertificateHandshake.buffer);
                    this.nextEpoch.verificationStream.push(payloadReader.buffer);

                    this.nextEpoch.serverPublicKey = certificate.publicKey.keyRaw;
                    this.nextEpoch.state = HandshakeState.ExpectingServerKeyExchange;
                    break;
                case HandshakeType.ServerKeyExchange:
                    if (this.nextEpoch.state !== HandshakeState.ExpectingServerKeyExchange) {
                        continue;
                    } else if (!this.nextEpoch.serverPublicKey) {
                        continue;
                    } else if (!this.nextEpoch.handshake) {
                        continue;
                    } else if (handshake.messageSequence !== 3) {
                        continue;
                    }

                    const sharedSecret = this.nextEpoch.handshake.verifyServerMessageAndGenerateSharedKey(payloadReader.buffer.slice(payloadReader.cursor), this.nextEpoch.serverPublicKey);
                    if (!sharedSecret) {
                        continue;
                    }

                    const randomSeed = Buffer.alloc(64);
                    this.nextEpoch.clientRandom.copy(randomSeed);
                    this.nextEpoch.serverRandom.copy(randomSeed, this.nextEpoch.clientRandom.byteLength);

                    const masterSecretSize = 48;
                    const masterSecret = Buffer.alloc(masterSecretSize);
                    expandSecret(masterSecret, sharedSecret, "master secert", randomSeed);

                    switch (this.nextEpoch.selectedCipherSuite) {
                        case CipherSuite.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256:
                            this.nextEpoch.recordProtection = new AesGcmRecordProtection(
                                masterSecret,
                                this.nextEpoch.serverRandom,
                                this.nextEpoch.clientRandom
                            );
                            break;
                    }

                    this.nextEpoch.state = HandshakeState.ExpectingServerHelloDone;
                    this.nextEpoch.masterSecret = masterSecret;

                    this.nextEpoch.verificationStream.push(reader.buffer.slice(startOfHandshake, startOfHandshake + 12 + handshake.length));
                    break;
                case HandshakeType.ServerHelloDone:
                    if (this.nextEpoch.state !== HandshakeState.ExpectingServerHelloDone) {
                        continue;
                    } else if (handshake.messageSequence !== 4) {
                        continue;
                    }

                    this.nextEpoch.state = HandshakeState.ExpectingChangeCipherSpec;
                    this.nextEpoch.verificationStream.push(reader.buffer.slice(startOfHandshake, startOfHandshake + 12 + handshake.length));

                    this.sendClientKeyExchangeFlight(false);
                    break;
                case HandshakeType.Finished:
                    if (this.nextEpoch.state !== HandshakeState.ExpectingFinished) {
                        continue;
                    } else if (handshake.messageSequence !== 7) {
                        continue;
                    }

                    if (Buffer.compare(reader.buffer, this.nextEpoch.serverVerification) !== 1) {
                        continue;
                    }

                    this.nextEpoch.epoch++;
                    this.nextEpoch.state = HandshakeState.Established;
                    this.nextEpoch.nextPacketResendTime = 0;
                    this.nextEpoch.serverVerification = Buffer.alloc(0);
                    this.nextEpoch.masterSecret = Buffer.alloc(32);

                    this.flushQueuedApplicationData();
                    break;
            }
        }
    }

    sendClientHello() {
        this.nextEpoch.verificationStream = [];

        const clientHello = new ClientHello(
            this.nextEpoch.clientRandom,
            this.nextEpoch.cookie,
            [ CipherSuite.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256 ],
            [
                new EllipticCurvesExtension([ NamedCurve.x25519 ])
            ]
        );

        const clientHelloWriter = HazelWriter.alloc(0);
        clientHelloWriter.write(clientHello);

        const handshake = new Handshake(
            HandshakeType.ClientHello,
            clientHelloWriter.size,
            0,
            0,
            clientHelloWriter.size
        );

        const plaintextLen = 12 + handshake.length;
        const outgoingRecord = new RecordHeader(
            ContentType.Handshake,
            this.epoch,
            this.currentEpoch.nextOutgoingSequence,
            this.currentEpoch.recordProtection.getEncryptedSize(plaintextLen)
        );
        this.currentEpoch.nextOutgoingSequence++;

        const packet = HazelWriter.alloc(RecordHeader.size + outgoingRecord.length);
        packet.write(outgoingRecord);
        packet.write(handshake);
        packet.bytes(clientHelloWriter.buffer);

        this.nextEpoch.verificationStream.push(packet.buffer.slice(RecordHeader.size, RecordHeader.size + 12 + handshake.length));

        const output = this.currentEpoch.recordProtection!.encryptClientPlaintext(
            packet.buffer.slice(RecordHeader.size, RecordHeader.size + plaintextLen),
            outgoingRecord
        );
        output.copy(packet.buffer, RecordHeader.size);

        this.nextEpoch.state = HandshakeState.ExpectingServerHello;
        this.nextEpoch.nextPacketResendTime = Date.now() + this.handshakeResendTimeout;

        this.socket?.send(packet.buffer);
    }

    sendClientKeyExchangeFlight(isRetrasmit: boolean) {
        const msgSize = 33;

        const keyExchangeHandshake = new Handshake(
            HandshakeType.ClientKeyExchange,
            msgSize,
            5,
            0,
            msgSize
        );

        const keyExchangeRecord = new RecordHeader(
            ContentType.Handshake,
            this.epoch,
            this.currentEpoch.nextOutgoingSequence,
            this.currentEpoch.recordProtection!.getEncryptedSize(12 + msgSize)
        );
        this.currentEpoch.nextOutgoingSequence++;

        const changeCipherSpecRecord = new RecordHeader(
            ContentType.ChangeCipherSpec,
            this.epoch,
            this.currentEpoch.nextOutgoingSequence,
            this.currentEpoch.recordProtection!.getEncryptedSize(1)
        );
        this.currentEpoch.nextOutgoingSequence++;

        const finishedHandshake = new Handshake(
            HandshakeType.Finished,
            12,
            6,
            0,
            12
        );

        const finishedRecord = new RecordHeader(
            ContentType.Handshake,
            this.nextEpoch.epoch,
            this.nextEpoch.nextOutgoingSequence,
            this.nextEpoch.recordProtection!.getEncryptedSize(12 + 12)
        );
        this.nextEpoch.nextOutgoingSequence++;

        const packetLen = 0 +
            RecordHeader.size + keyExchangeRecord.length +
            RecordHeader.size + changeCipherSpecRecord.length +
            RecordHeader.size + finishedRecord.length;

        const writer = HazelWriter.alloc(packetLen);
        writer.write(keyExchangeRecord);
        writer.write(keyExchangeHandshake);
        const keyExchangeEncode = Buffer.alloc(33);
        this.nextEpoch.handshake?.encodeClientKeyExchangeMessage(keyExchangeEncode);
        writer.bytes(keyExchangeEncode);

        const startOfChangeCipherSpecRecord = writer.cursor;
        writer.write(changeCipherSpecRecord);
        writer.write(ChangeCipherSpec);

        const startOfFinishedRecord = writer.cursor;
        writer.write(finishedRecord);
        writer.write(finishedHandshake);

        if (!isRetrasmit) {
            this.nextEpoch.verificationStream.push(
                writer.buffer.slice(RecordHeader.size, RecordHeader.size + 12 + keyExchangeHandshake.length)
            );
        }

        const handshakeHash = crypto.createHash("sha256").update(Buffer.concat(this.nextEpoch.verificationStream)).digest();
        expandSecret(this.nextEpoch.serverVerification, this.nextEpoch.masterSecret, "server finished", handshakeHash);
        const expandKeyOutput = Buffer.alloc(12);
        expandSecret(expandKeyOutput, this.nextEpoch.masterSecret, "client finished", handshakeHash);
        writer.bytes(expandKeyOutput);

        const keyExchangeEncrypt = this.currentEpoch.recordProtection.encryptClientPlaintext(
            writer.buffer.slice(RecordHeader.size, RecordHeader.size + keyExchangeHandshake.length),
            keyExchangeRecord
        );
        keyExchangeEncrypt.copy(writer.buffer, RecordHeader.size, 0, keyExchangeEncrypt.byteLength);

        const changeCipherSpecEncrypt = this.currentEpoch.recordProtection!.encryptClientPlaintext(
            writer.buffer.slice(startOfChangeCipherSpecRecord + RecordHeader.size, startOfChangeCipherSpecRecord + RecordHeader.size + 1),
            changeCipherSpecRecord
        );
        changeCipherSpecEncrypt.copy(writer.buffer, startOfChangeCipherSpecRecord + RecordHeader.size);

        const finishedEncrypt = this.nextEpoch.recordProtection!.encryptClientPlaintext(
            writer.buffer.slice(startOfFinishedRecord + RecordHeader.size, startOfFinishedRecord + RecordHeader.size + 12 + finishedHandshake.length),
            finishedRecord
        );
        finishedEncrypt.copy(writer.buffer, startOfFinishedRecord + RecordHeader.size);

        this.nextEpoch.state = HandshakeState.ExpectingChangeCipherSpec;
        this.nextEpoch.nextPacketResendTime = Date.now() + this.handshakeResendTimeout;
        this.socket?.send(writer.buffer);
    }
}

//------------------------------------------------------------------------------------------------------------------------------------------

export class RoombaV2 extends EventEmitter {
  public roomba?: Local;
  private timeout?: NodeJS.Timeout;
  private keepAlive = this.refreshInterval === -1;

  constructor(private readonly blid: string, private readonly password: string, private readonly ip: string,
    private readonly refreshInterval: number, private readonly log: Logger, private readonly logPrefix: string) {
    super();
    process.env.ROBOT_CIPHERS = 'AES128-SHA256';
  }

  connect(): Promise<Local> {
    if (this.timeout) {
      clearTimeout(this.timeout);
      this.timeout = undefined;
    }
    return new Promise((resolve, reject) => {
      if (this.roomba) {
        resolve(this.roomba);
      } else {
        this.log.debug(this.logPrefix, 'Connecting...');
        this.roomba = new Local(this.blid, this.password, this.ip, 2);
        this.roomba.on('offline', () => {
          this.log.debug(this.logPrefix, 'Offline');
          this.roomba.end();
          reject('Roomba Offline');
        }).on('close', () => {
          this.log.debug(this.logPrefix, 'Disconnected');
          this.roomba = undefined;
        }).on('connect', () => {
          this.log.debug(this.logPrefix, 'Connected');
          resolve(this.roomba);
        }).on('state', (state) => {
          this.emit('update', Object.assign(state, state.cleanMissionStatus, state.bin));
        });
      }
    });
  }

  disconnect(roomba: Local) {
    if (this.timeout) {
      clearTimeout(this.timeout);
      this.timeout = undefined;
    }
    if (!this.keepAlive) {
      this.timeout = setTimeout(() => {
        this.log.debug(this.logPrefix, 'Disconnecting...');
        roomba.end();
      }, this.refreshInterval || 5000);

    }
  }

  end() {
    this.connect().then((roomba) => {
      roomba.end();
    });
  }

  clean() {
    this.connect().then(async (roomba) => {
      await roomba.clean();
      this.disconnect(roomba);
    });
  }

  pause() {
    this.connect().then(async (roomba) => {
      await roomba.pause();
      this.disconnect(roomba);
    });
  }

  stop() {
    this.connect().then(async (roomba) => {
      await roomba.stop();
      this.disconnect(roomba);
    });
  }

  resume() {
    this.connect().then(async (roomba) => {
      await roomba.resume();
      this.disconnect(roomba);
    });
  }

  dock() {
    this.connect().then(async (roomba) => {
      await roomba.dock();
      this.disconnect(roomba);
    });
  }

  find() {
    this.connect().then(async (roomba) => {
      await roomba.find();
      this.disconnect(roomba);
    });
  }

  async getMission(): Promise<MissionV2> {
    return new Promise((resolve, reject) => {
      setTimeout(() =>{
        reject('Operation Timed out');
      }, 3000);
      this.connect().then(async (roomba) => {
        roomba.getRobotState(['cleanMissionStatus', 'bin', 'batPct'])
          .then(state => resolve(Object.assign(state, state.cleanMissionStatus, state.bin)))
          .catch(err => reject(err));
        this.disconnect(roomba);
      }).catch(err => reject(err));
    });
  }
}

export class RoombaV3 extends EventEmitter {
  public roomba?: Local;
  private timeout?: NodeJS.Timeout;
  private keepAlive = this.refreshInterval === -1;

  constructor(private readonly blid: string, private readonly password: string, private readonly ip: string, private readonly sku: string,
    private readonly refreshInterval: number, private readonly log: Logger, private readonly logPrefix: string) {
    super();
    process.env.ROBOT_CIPHERS = this.sku.startsWith('j') ? 'TLS_AES_256_GCM_SHA384' : 'AES128-SHA256';
  }

  connect(): Promise<Local> {
    if (this.timeout) {
      clearTimeout(this.timeout);
      this.timeout = undefined;
    }
    return new Promise((resolve, reject) => {
      if (this.roomba) {
        resolve(this.roomba);
      } else {
        this.log.debug('Connecting...');
        this.roomba = new Local(this.blid, this.password, this.ip, 2);
        this.roomba.on('offline', () => {
          this.log.debug(this.logPrefix, 'Offline');
          this.roomba.end();
          reject('Roomba Offline');
        }).on('close', () => {
          this.log.debug(this.logPrefix, 'Disconnected');
          this.roomba = undefined;
        }).on('connect', () => {
          this.log.debug(this.logPrefix, 'Connected');
          resolve(this.roomba);
        }).on('update', (state) => {
          this.emit('state', Object.assign(state, state.cleanMissionStatus, state.bin));
        });
      }
    });
  }

  disconnect(roomba: Local) {
    if (this.timeout) {
      clearTimeout(this.timeout);
      this.timeout = undefined;
    }
    if (!this.keepAlive) {
      this.timeout = setTimeout(() => {
        this.log.debug(this.logPrefix, 'Disconnecting...');
        roomba.end();
      }, this.refreshInterval || 5000);
    }
  }

  end() {
    this.connect().then((roomba) => {
      roomba.end();
    });
  }

  clean() {
    this.connect().then(async (roomba) => {
      await roomba.clean();
      this.disconnect(roomba);
    });
  }

  cleanRoom(map) {
    map.ordered = 1;
    this.connect().then(async (roomba) => {
      await roomba.cleanRoom(map);
      this.disconnect(roomba);
    });
  }

  pause() {
    this.connect().then(async (roomba) => {
      await roomba.pause();
      this.disconnect(roomba);
    });
  }

  stop() {
    this.connect().then(async (roomba) => {
      await roomba.stop();
      this.disconnect(roomba);
    });
  }

  resume() {
    this.connect().then(async (roomba) => {
      await roomba.resume();
      this.disconnect(roomba);
    });
  }

  dock() {
    this.connect().then(async (roomba) => {
      await roomba.dock();
      this.disconnect(roomba);
    });
  }

  find() {
    this.connect().then(async (roomba) => {
      await roomba.find();
      this.disconnect(roomba);
    });
  }

  async getMission(): Promise<MissionV3> {
    return new Promise((resolve, reject) => {
      setTimeout(() =>{
        reject('Operation Timed out');
      }, 3000);
      this.connect().then(async (roomba) => {
        roomba.getRobotState(['cleanMissionStatus', 'bin', 'batPct'])
          .then(state => resolve(Object.assign(state, state.cleanMissionStatus, state.bin)))
          .catch(err => reject(err));
        this.disconnect(roomba);
      }).catch(err => reject(err));
    });
  }
}

suite('StripeDaemon', () => {
  let daemonProcessStub: execa.ExecaChildProcess;
  let sandbox: sinon.SinonSandbox;

  const stripeClient = <Partial<StripeClient>>{
    getCLIPath: () => Promise.resolve('/path/to/cli'),
    promptUpdateForDaemon: () => {},
  };

  // Get an instance of StripeDaemon with the mocked execa module
  const getStripeDaemonWithExecaProxy = (
    stdout: string,
    stripeClient: StripeClient,
  ): StripeDaemon => {
    daemonProcessStub = <execa.ExecaChildProcess<string>>new EventEmitter();
    daemonProcessStub.stdin = new Writable({write: () => {}});
    daemonProcessStub.stdout = new Readable({
      read() {
        this.push(stdout, 'utf8');
        this.push(null); // nothing left to read
      },
    });
    daemonProcessStub.stderr = new Readable({read: () => {}});
    daemonProcessStub.kill = () => true;

    const module = setupProxies({execa: () => daemonProcessStub});
    return new module.StripeDaemon(stripeClient);
  };

  setup(() => {
    sandbox = sinon.createSandbox();
  });

  teardown(() => {
    sandbox.restore();
  });

  suite('setupClient', () => {
    test('returns a new client that connects to the daemon address', async () => {
      const stripeDaemon = getStripeDaemonWithExecaProxy(
        '{"host": "::1", "port": 12345}',
        <any>stripeClient,
      );

      const constructorStub = sandbox.stub();
      Object.setPrototypeOf(StripeCLIClient, constructorStub);

      await stripeDaemon.setupClient();
      console.log(constructorStub.args[0][2].channelOverride.options['grpc.primary_user_agent']);
      assert.strictEqual(constructorStub.args[0][0], '[::1]:12345');
    });

    test('sends correct channel options', async () => {
      const stripeDaemon = getStripeDaemonWithExecaProxy(
        '{"host": "::1", "port": 12345}',
        <any>stripeClient,
      );

      const constructorStub = sandbox.stub();
      Object.setPrototypeOf(StripeCLIClient, constructorStub);

      // mock out extensionId
      const mockExtension = <vscode.Extension<any>>{
        id: 'my-extension',
        packageJSON: {version: 1},
      };
      sandbox.stub(vscode.extensions, 'getExtension').returns(mockExtension);

      await stripeDaemon.setupClient();

      const userAgent =
        constructorStub.args[0][2].channelOverride.options['grpc.primary_user_agent'];

      // Note I could not mock out the module that's used within the utils class so we are just asserting for a startsWith
      assert.strictEqual(userAgent.startsWith('my-extension/1 vscode/'), true);
    });

    test('rejects with SyntaxError when stdout is invalid json', async () => {
      const stripeDaemon = getStripeDaemonWithExecaProxy(
        'unexpected string from stripe daemon',
        <any>stripeClient,
      );
      await assert.rejects(stripeDaemon.setupClient(), SyntaxError);
    });

    test('rejects with MalformedConfigError when stdout is valid json but not a daemon config', async () => {
      const stripeDaemon = getStripeDaemonWithExecaProxy('{"foo": "bar"}', <any>stripeClient);
      await assert.rejects(stripeDaemon.setupClient(), {
        name: 'MalformedConfigError',
        message: 'Received malformed config from stripe daemon: {"foo":"bar"}',
      });
    });

    test('rejects with NoDaemonCommandError when daemon command does not exist', async () => {
      const stripeDaemon = getStripeDaemonWithExecaProxy(
        'Unknown command "daemon" for "stripe".',
        <any>stripeClient,
      );
      await assert.rejects(stripeDaemon.setupClient(), {
        name: 'NoDaemonCommandError',
        message: 'Daemon is not available with this CLI version',
      });
    });
  });
});

suite('stripeEventsView', () => {
  let sandbox: sinon.SinonSandbox;

  const workspaceState = new TestMemento();
  const extensionContext = {...mocks.extensionContextMock, workspaceState: workspaceState};

  const stripeClient = <Partial<StripeClient>>{
    getCLIPath: () => Promise.resolve('/path/to/cli'),
    promptUpdateForDaemon: () => {},
    promptLogin: () => {},
  };

  const stripeDaemon = <Partial<StripeDaemon>>{
    setupClient: () => {},
  };

  let listenStream: grpc.ClientReadableStream<ListenResponse>;
  let daemonClient: Partial<StripeCLIClient>;

  setup(() => {
    sandbox = sinon.createSandbox();

    listenStream = <grpc.ClientReadableStream<ListenResponse>>new EventEmitter();
    listenStream.cancel = () => {};
    listenStream.destroy = () => {};

    daemonClient = {
      listen: () => listenStream,
    };
  });

  teardown(() => {
    sandbox.restore();
  });

  suite('startStreaming', () => {
    setup(() => {
      sandbox.stub(stripeDaemon, 'setupClient').resolves(daemonClient);
    });

    suite('state transitions', () => {
      test('renders loading state when stream is LOADING', async () => {
        const stripeEventsView = new StripeEventsViewProvider(
          <any>stripeClient,
          <any>stripeDaemon,
          extensionContext,
        );
        await stripeEventsView.startStreaming();

        // Make sure we start in the idle state
        const doneResponse = new ListenResponse();
        doneResponse.setState(ListenResponse.State.STATE_DONE);
        listenStream.emit('data', doneResponse);

        const loadingResponse = new ListenResponse();
        loadingResponse.setState(ListenResponse.State.STATE_LOADING);

        listenStream.emit('data', loadingResponse);

        const treeItems = await stripeEventsView.buildTree(); // Simulate view refresh

        const labels = treeItems.map(({label}) => label);
        const expectedLabel = 'Starting streaming events ...';
        assert.strictEqual(
          labels.includes(expectedLabel),
          true,
          `Expected [${labels.toString()}] to contain ${expectedLabel}`,
        );
      });

      test('renders loading state when stream is RECONNECTING', async () => {
        const stripeEventsView = new StripeEventsViewProvider(
          <any>stripeClient,
          <any>stripeDaemon,
          extensionContext,
        );
        await stripeEventsView.startStreaming();

        // Make sure we start in the idle state
        const doneResponse = new ListenResponse();
        doneResponse.setState(ListenResponse.State.STATE_DONE);
        listenStream.emit('data', doneResponse);

        const reconnectingResponse = new ListenResponse();
        reconnectingResponse.setState(ListenResponse.State.STATE_RECONNECTING);

        listenStream.emit('data', reconnectingResponse);

        const treeItems = await stripeEventsView.buildTree(); // Simulate view refresh

        const labels = treeItems.map(({label}) => label);
        const expectedLabel = 'Starting streaming events ...';
        assert.strictEqual(
          labels.includes(expectedLabel),
          true,
          `Expected [${labels.toString()}] to contain ${expectedLabel}`,
        );
      });

      test('renders ready state when stream is READY', async () => {
        const stripeEventsView = new StripeEventsViewProvider(
          <any>stripeClient,
          <any>stripeDaemon,
          extensionContext,
        );
        await stripeEventsView.startStreaming();

        // Make sure we start in the streaming state
        const doneResponse = new ListenResponse();
        doneResponse.setState(ListenResponse.State.STATE_DONE);
        listenStream.emit('data', doneResponse);

        const readyResponse = new ListenResponse();
        readyResponse.setState(ListenResponse.State.STATE_READY);

        listenStream.emit('data', readyResponse);

        const treeItems = await stripeEventsView.buildTree(); // Simulate view refresh

        const labels = treeItems.map(({label}) => label);
        const expectedLabel = 'Stop streaming events';
        assert.strictEqual(
          labels.includes(expectedLabel),
          true,
          `Expected [${labels.toString()}] to contain ${expectedLabel}`,
        );
      });

      test('renders idle state when stream is DONE', async () => {
        const stripeEventsView = new StripeEventsViewProvider(
          <any>stripeClient,
          <any>stripeDaemon,
          extensionContext,
        );
        await stripeEventsView.startStreaming();

        // Make sure we start in the streaming state
        const readyResponse = new ListenResponse();
        readyResponse.setState(ListenResponse.State.STATE_READY);
        listenStream.emit('data', readyResponse);

        const doneResponse = new ListenResponse();
        doneResponse.setState(ListenResponse.State.STATE_DONE);

        listenStream.emit('data', doneResponse);

        const treeItems = await stripeEventsView.buildTree(); // Simulate view refresh

        const labels = treeItems.map(({label}) => label);
        const expectedLabel = 'Start streaming events';
        assert.strictEqual(
          labels.includes(expectedLabel),
          true,
          `Expected [${labels.toString()}] to contain ${expectedLabel}`,
        );
      });

      test('renders idle state when stream is receives unknown state', async () => {
        const stripeEventsView = new StripeEventsViewProvider(
          <any>stripeClient,
          <any>stripeDaemon,
          extensionContext,
        );
        await stripeEventsView.startStreaming();

        // Make sure we start in the ready state
        const readyResponse = new ListenResponse();
        readyResponse.setState(ListenResponse.State.STATE_READY);
        listenStream.emit('data', readyResponse);

        const unknownStateResponse = new ListenResponse();
        unknownStateResponse.setState(<any>-1); // This should be impossible

        listenStream.emit('data', unknownStateResponse);

        const treeItems = await stripeEventsView.buildTree(); // Simulate view refresh

        const labels = treeItems.map(({label}) => label);
        const expectedLabel = 'Start streaming events';
        assert.strictEqual(
          labels.includes(expectedLabel),
          true,
          `Expected [${labels.toString()}] to contain ${expectedLabel}`,
        );
      });
    });

    test('creates tree items from stream', async () => {
      const stripeEventsView = new StripeEventsViewProvider(
        <any>stripeClient,
        <any>stripeDaemon,
        extensionContext,
      );
      await stripeEventsView.startStreaming();

      // Mock ready response
      const readyResponse = new ListenResponse();
      readyResponse.setState(ListenResponse.State.STATE_READY);

      listenStream.emit('data', readyResponse);

      // Mock event response
      const stripeEvent = new StripeEvent();
      stripeEvent.setType('customer.created');
      stripeEvent.setId('evt_123');

      const response = new ListenResponse();
      response.setStripeEvent(stripeEvent);

      listenStream.emit('data', response);

      const treeItems = await stripeEventsView.buildTree(); // Simulate view refresh

      const recentEvents = treeItems.find(({label}) => label === 'Recent events');
      const labels = recentEvents?.children?.map(({label}) => label);
      const expectedLabel = 'customer.created';
      assert.strictEqual(
        labels?.includes(expectedLabel),
        true,
        `Expected [${labels?.toString()}] to contain ${expectedLabel}`,
      );
    });
  });

  suite('stopStreaming', () => {
    setup(() => {
      sandbox.stub(stripeDaemon, 'setupClient').resolves(daemonClient);
    });

    test('stops streaming', async () => {
      // Simulate a stream in progress
      const stripeEventsView = new StripeEventsViewProvider(
        <any>stripeClient,
        <any>stripeDaemon,
        extensionContext,
      );
      await stripeEventsView.startStreaming();

      const readyResponse = new ListenResponse();
      readyResponse.setState(ListenResponse.State.STATE_READY);

      listenStream.emit('data', readyResponse);

      stripeEventsView.stopStreaming();

      const treeItems = await stripeEventsView.buildTree(); // Simulate view refresh

      const labels = treeItems.map(({label}) => label);
      const expectedLabel = 'Start streaming events';
      assert.strictEqual(
        labels.includes(expectedLabel),
        true,
        `Expected [${labels.toString()}] to contain ${expectedLabel}`,
      );
    });
  });

  suite('error', () => {
    test('prompts upgrade when no daemon command', async () => {
      sandbox.stub(stripeDaemon, 'setupClient').throws(new NoDaemonCommandError());

      const promptUpdateForDaemonSpy = sandbox.spy(stripeClient, 'promptUpdateForDaemon');

      const stripeEventsView = new StripeEventsViewProvider(
        <any>stripeClient,
        <any>stripeDaemon,
        extensionContext,
      );

      await stripeEventsView.startStreaming();

      assert.strictEqual(promptUpdateForDaemonSpy.calledOnce, true);
    });

    test('shows gRPC error message', async () => {
      sandbox.stub(stripeDaemon, 'setupClient').resolves(daemonClient);

      const showErrorMessageSpy = sandbox.spy(vscode.window, 'showErrorMessage');

      const stripeEventsView = new StripeEventsViewProvider(
        <any>stripeClient,
        <any>stripeDaemon,
        extensionContext,
      );

      await stripeEventsView.startStreaming();

      listenStream.emit('error', <Partial<grpc.ServiceError>>{
        code: grpc.status.UNKNOWN,
        details: 'unknown error',
      });

      assert.strictEqual(showErrorMessageSpy.callCount, 1);

      assert.strictEqual(showErrorMessageSpy.args[0][0], 'unknown error');
    });

    test('prompts login when UNAUTHENTICATED', async () => {
      sandbox.stub(stripeDaemon, 'setupClient').resolves(daemonClient);

      const promptLoginSpy = sandbox.spy(stripeClient, 'promptLogin');

      const stripeEventsView = new StripeEventsViewProvider(
        <any>stripeClient,
        <any>stripeDaemon,
        extensionContext,
      );

      await stripeEventsView.startStreaming();

      listenStream.emit('error', <Partial<grpc.ServiceError>>{code: grpc.status.UNAUTHENTICATED});

      assert.strictEqual(promptLoginSpy.callCount, 1);
    });

    test('silently handle CANCELLED error', async () => {
      sandbox.stub(stripeDaemon, 'setupClient').resolves(daemonClient);

      const showErrorMessageSpy = sandbox.spy(vscode.window, 'showErrorMessage');

      const stripeEventsView = new StripeEventsViewProvider(
        <any>stripeClient,
        <any>stripeDaemon,
        extensionContext,
      );

      await stripeEventsView.startStreaming();

      listenStream.emit('error', <Partial<grpc.ServiceError>>{code: grpc.status.CANCELLED});

      assert.strictEqual(showErrorMessageSpy.callCount, 0);
    });
  });
});

suite('stripeLogsView', () => {
  let sandbox: sinon.SinonSandbox;

  const workspaceState = new TestMemento();
  const extensionContext = {...mocks.extensionContextMock, workspaceState: workspaceState};

  const stripeClient = <Partial<StripeClient>>{
    getCLIPath: () => Promise.resolve('/path/to/cli'),
    promptUpdateForDaemon: () => {},
    promptLogin: () => {},
  };

  const stripeDaemon = <Partial<StripeDaemon>>{
    setupClient: () => {},
  };

  let logsTailStream: grpc.ClientReadableStream<LogsTailResponse>;
  let daemonClient: Partial<StripeCLIClient>;

  setup(() => {
    sandbox = sinon.createSandbox();

    logsTailStream = <grpc.ClientReadableStream<LogsTailResponse>>new EventEmitter();
    logsTailStream.cancel = () => {};
    logsTailStream.destroy = () => {};

    daemonClient = {
      logsTail: () => logsTailStream,
    };
  });

  teardown(() => {
    sandbox.restore();
  });

  suite('startStreaming', () => {
    setup(() => {
      sandbox.stub(stripeDaemon, 'setupClient').resolves(daemonClient);
    });

    suite('state transitions', () => {
      test('renders loading state when stream is LOADING', async () => {
        const stripeLogsView = new StripeLogsViewProvider(
          <any>stripeClient,
          <any>stripeDaemon,
          extensionContext,
        );
        await stripeLogsView.startStreaming();

        // Make sure we start in the idle state
        const doneResponse = new LogsTailResponse();
        doneResponse.setState(LogsTailResponse.State.STATE_DONE);
        logsTailStream.emit('data', doneResponse);

        const loadingResponse = new LogsTailResponse();
        loadingResponse.setState(LogsTailResponse.State.STATE_LOADING);

        logsTailStream.emit('data', loadingResponse);

        const treeItems = await stripeLogsView.buildTree(); // Simulate view refresh

        const labels = treeItems.map(({label}) => label);
        const expectedLabel = 'Starting streaming API logs ...';
        assert.strictEqual(
          labels.includes(expectedLabel),
          true,
          `Expected [${labels.toString()}] to contain ${expectedLabel}`,
        );
      });

      test('renders loading state when stream is RECONNECTING', async () => {
        const stripeLogsView = new StripeLogsViewProvider(
          <any>stripeClient,
          <any>stripeDaemon,
          extensionContext,
        );
        await stripeLogsView.startStreaming();

        // Make sure we start in the idle state
        const doneResponse = new LogsTailResponse();
        doneResponse.setState(LogsTailResponse.State.STATE_DONE);
        logsTailStream.emit('data', doneResponse);

        const reconnectingResponse = new LogsTailResponse();
        reconnectingResponse.setState(LogsTailResponse.State.STATE_RECONNECTING);

        logsTailStream.emit('data', reconnectingResponse);

        const treeItems = await stripeLogsView.buildTree(); // Simulate view refresh

        const labels = treeItems.map(({label}) => label);
        const expectedLabel = 'Starting streaming API logs ...';
        assert.strictEqual(
          labels.includes(expectedLabel),
          true,
          `Expected [${labels.toString()}] to contain ${expectedLabel}`,
        );
      });

      test('renders ready state when stream is READY', async () => {
        const stripeLogsView = new StripeLogsViewProvider(
          <any>stripeClient,
          <any>stripeDaemon,
          extensionContext,
        );
        await stripeLogsView.startStreaming();

        // Make sure we start in the streaming state
        const doneResponse = new LogsTailResponse();
        doneResponse.setState(LogsTailResponse.State.STATE_DONE);
        logsTailStream.emit('data', doneResponse);

        const readyResponse = new LogsTailResponse();
        readyResponse.setState(LogsTailResponse.State.STATE_READY);

        logsTailStream.emit('data', readyResponse);

        const treeItems = await stripeLogsView.buildTree(); // Simulate view refresh

        const labels = treeItems.map(({label}) => label);
        const expectedLabel = 'Stop streaming API logs';
        assert.strictEqual(
          labels.includes(expectedLabel),
          true,
          `Expected [${labels.toString()}] to contain ${expectedLabel}`,
        );
      });

      test('renders idle state when stream is DONE', async () => {
        const stripeLogsView = new StripeLogsViewProvider(
          <any>stripeClient,
          <any>stripeDaemon,
          extensionContext,
        );
        await stripeLogsView.startStreaming();

        // Make sure we start in the streaming state
        const readyResponse = new LogsTailResponse();
        readyResponse.setState(LogsTailResponse.State.STATE_READY);
        logsTailStream.emit('data', readyResponse);

        const doneResponse = new LogsTailResponse();
        doneResponse.setState(LogsTailResponse.State.STATE_DONE);

        logsTailStream.emit('data', doneResponse);

        const treeItems = await stripeLogsView.buildTree(); // Simulate view refresh

        const labels = treeItems.map(({label}) => label);
        const expectedLabel = 'Start streaming API logs';
        assert.strictEqual(
          labels.includes(expectedLabel),
          true,
          `Expected [${labels.toString()}] to contain ${expectedLabel}`,
        );
      });

      test('renders idle state when stream is receives unknown state', async () => {
        const stripeLogsView = new StripeLogsViewProvider(
          <any>stripeClient,
          <any>stripeDaemon,
          extensionContext,
        );
        await stripeLogsView.startStreaming();

        // Make sure we start in the ready state
        const readyResponse = new LogsTailResponse();
        readyResponse.setState(LogsTailResponse.State.STATE_READY);
        logsTailStream.emit('data', readyResponse);

        const unknownStateResponse = new LogsTailResponse();
        unknownStateResponse.setState(<any>-1); // This should be impossible

        logsTailStream.emit('data', unknownStateResponse);

        const treeItems = await stripeLogsView.buildTree(); // Simulate view refresh

        const labels = treeItems.map(({label}) => label);
        const expectedLabel = 'Start streaming API logs';
        assert.strictEqual(
          labels.includes(expectedLabel),
          true,
          `Expected [${labels.toString()}] to contain ${expectedLabel}`,
        );
      });
    });

    test('creates tree items from stream', async () => {
      const stripeLogsView = new StripeLogsViewProvider(
        <any>stripeClient,
        <any>stripeDaemon,
        extensionContext,
      );
      await stripeLogsView.startStreaming();

      // Mock ready response
      const readyResponse = new LogsTailResponse();
      readyResponse.setState(LogsTailResponse.State.STATE_READY);

      logsTailStream.emit('data', readyResponse);

      // Mock log response
      const log = new LogsTailResponse.Log();
      log.setStatus(200);
      log.setMethod('POST');
      log.setUrl('/v1/customers');
      log.setRequestId('req_123');
      log.setCreatedAt(12345);

      const response = new LogsTailResponse();
      response.setLog(log);

      logsTailStream.emit('data', response);

      const treeItems = await stripeLogsView.buildTree(); // Simulate view refresh

      const recentLogs = treeItems.find(({label}) => label === 'Recent logs');
      const labels = recentLogs?.children?.map(({label}) => label);
      const expectedLabel = '[200] POST /v1/customers [req_123]';
      assert.strictEqual(
        labels?.includes(expectedLabel),
        true,
        `Expected [${labels?.toString()}] to contain ${expectedLabel}`,
      );
    });
  });

  suite('stopStreaming', () => {
    setup(() => {
      sandbox.stub(stripeDaemon, 'setupClient').resolves(daemonClient);
    });

    test('stops streaming', async () => {
      // Simulate a stream in progress
      const stripeLogsView = new StripeLogsViewProvider(
        <any>stripeClient,
        <any>stripeDaemon,
        extensionContext,
      );
      await stripeLogsView.startStreaming();

      const readyResponse = new LogsTailResponse();
      readyResponse.setState(LogsTailResponse.State.STATE_READY);

      logsTailStream.emit('data', readyResponse);

      stripeLogsView.stopStreaming();

      const treeItems = await stripeLogsView.buildTree(); // Simulate view refresh

      const labels = treeItems.map(({label}) => label);
      const expectedLabel = 'Start streaming API logs';
      assert.strictEqual(
        labels.includes(expectedLabel),
        true,
        `Expected [${labels.toString()}] to contain ${expectedLabel}`,
      );
    });
  });

  suite('error', () => {
    test('prompts upgrade when no daemon command', async () => {
      sandbox.stub(stripeDaemon, 'setupClient').throws(new NoDaemonCommandError());

      const promptUpdateForDaemonSpy = sandbox.spy(stripeClient, 'promptUpdateForDaemon');

      const stripeLogsView = new StripeLogsViewProvider(
        <any>stripeClient,
        <any>stripeDaemon,
        extensionContext,
      );

      await stripeLogsView.startStreaming();

      assert.strictEqual(promptUpdateForDaemonSpy.calledOnce, true);
    });

    test('shows gRPC error message', async () => {
      sandbox.stub(stripeDaemon, 'setupClient').resolves(daemonClient);

      const showErrorMessageSpy = sandbox.spy(vscode.window, 'showErrorMessage');

      const stripeLogsView = new StripeLogsViewProvider(
        <any>stripeClient,
        <any>stripeDaemon,
        extensionContext,
      );

      await stripeLogsView.startStreaming();

      logsTailStream.emit('error', <Partial<grpc.ServiceError>>{
        code: grpc.status.UNKNOWN,
        details: 'unknown error',
      });

      assert.strictEqual(showErrorMessageSpy.callCount, 1);

      assert.strictEqual(showErrorMessageSpy.args[0][0], 'unknown error');
    });

    test('prompts login when UNAUTHENTICATED', async () => {
      sandbox.stub(stripeDaemon, 'setupClient').resolves(daemonClient);

      const promptLoginSpy = sandbox.spy(stripeClient, 'promptLogin');

      const stripeLogsView = new StripeLogsViewProvider(
        <any>stripeClient,
        <any>stripeDaemon,
        extensionContext,
      );

      await stripeLogsView.startStreaming();

      logsTailStream.emit('error', <Partial<grpc.ServiceError>>{code: grpc.status.UNAUTHENTICATED});

      assert.strictEqual(promptLoginSpy.callCount, 1);
    });

    test('silently handle CANCELLED error', async () => {
      sandbox.stub(stripeDaemon, 'setupClient').resolves(daemonClient);

      const showErrorMessageSpy = sandbox.spy(vscode.window, 'showErrorMessage');

      const stripeLogsView = new StripeLogsViewProvider(
        <any>stripeClient,
        <any>stripeDaemon,
        extensionContext,
      );

      await stripeLogsView.startStreaming();

      logsTailStream.emit('error', <Partial<grpc.ServiceError>>{code: grpc.status.CANCELLED});

      assert.strictEqual(showErrorMessageSpy.callCount, 0);
    });
  });
});