Dart/Flutter mocking stream and with await for

Question

I have this kind of code

void main() {
  runApp(Provider(
    create: (_) => Repo(),
    child: MyApp(),
  ));
}

class MyApp extends StatefulWidget {
  @override
  State<MyApp> createState() => _MyAppState();
}

class _MyAppState extends State<MyApp> {
  WebSocket? _ws;
  var value = "";
  Timer? _timer;

  void _connect() async {
    try {
      print("connecting...");
      _ws = await context.read<Repo>().websocket.timeout(Duration(seconds: 10));
      print("connected");

        await for (var s in _ws!) {
        setState(() {
          value = s;
        });
      }
    } finally {
      print("disconnected");
      _timer = Timer(Duration(seconds: 10),  _connect);
    }
  }

  @override
  void initState() {
    _connect();
    super.initState();
  }

  @override
  void dispose() {
    _timer?.cancel();
    _ws?.close();
    super.dispose();
  }

  @override
  Widget build(BuildContext context) {
    return MaterialApp(
      home: Scaffold(
        body: Text(value),
      ),
    );
  }
}

class Repo {
  Future<WebSocket> get websocket => WebSocket.connect('ws://...');
}

This code works. I would like to test the automatic reconnection when websocket connection is closed (with mocktail to mock). I discovered that stream.listen(...) is called under the hood when using await for syntax. But I'm unable to get out of the loop during test. This is what I tried.

import 'dart:async';
import 'dart:io';

import 'package:flutter_playground/main.dart';
import 'package:flutter_test/flutter_test.dart';
import 'package:mocktail/mocktail.dart';
import 'package:provider/provider.dart';

class MockWS extends Mock implements WebSocket {}

class MockRepo extends Mock implements Repo {}

void main() {
  testWidgets('auto reconnect', (WidgetTester tester) async {
    final ws = MockWS();
    final repo = MockRepo();
    final ctrl = StreamController();
    when(() => repo.websocket).thenAnswer((_) async => ws);
    when(() => ws.listen(
          any(),
          onDone: any(named: 'onDone'),
          onError: any(named: 'onError'),
          cancelOnError: any(named: 'cancelOnError'),
        )).thenAnswer((invocation) {
      return ctrl.stream.listen(
        invocation.positionalArguments[0],
        onDone: invocation.namedArguments['onDone'],
        onError: invocation.namedArguments['onError'],
        cancelOnError: invocation.namedArguments['cancelOnError'],
      );
    });
    when(() => ws.close()).thenAnswer((invocation) async => null);
    await tester.pumpWidget(Provider<Repo>.value(
      value: repo,
      child: MyApp(),
    ));

    // Try to close the stream
    await ctrl.close();
    await tester.pump(Duration(seconds: 10));
    verify(() => repo.websocket).called(2);
  });
}

But the test fails, it does not get out the await for loop. How to emulate a loop exit ?