Widget Testing for AI Features in Flutter

An AI feature is a state machine wearing a widget. Behind a single “Send” button you are juggling a loading indicator, tokens that stream in one by one, and failures that arrive seconds after the user gave up hope. The worst part is that a live model is a terrible test double: nondeterministic, rate-limited, and expensive. This post shows you how to lock that state machine down with Flutter’s widget testing toolkit so your chat UI degrades gracefully, streams correctly, and fails loudly.

In this tutorial, you will learn how to:

• Structure widget tests with testWidgets, WidgetTester, and finders from the flutter_test package
• Mock an LLM client with a hand-written fake and with Mockito
• Test loading, streaming, and error states of an AI chat screen without a network
• Assert on partial streamed output using pump() and pumpAndSettle()
• Capture golden images for your AI UI and manage their baselines

Key technologies: flutter_test, testWidgets, WidgetTester, Finder, Mockito (or mocktail), Stream, and golden file testing.

Prerequisites

• Flutter SDK 3.x with the flutter_test package already in your dev_dependencies (it ships with every new project)
• A Flutter app with an abstract LLM service you can inject into your widgets
• Optional: mockito and build_runner, or mocktail if you prefer no code generation

The widget under test: a streaming chat screen

Before testing anything you need a seam. Abstract the model behind an interface so your widget never depends on a concrete SDK:

abstract class LlmClient {
  Stream<String> streamChat(String prompt);
}

A Stream<String> is the right contract for streaming tokens: each yield is one chunk the model produced, and the stream’s done event signals the end of the response. Here is a chat screen that drives it:

class ChatScreen extends StatefulWidget {
  const ChatScreen({super.key, required this.client});

  final LlmClient client;

  @override
  State<ChatScreen> createState() => _ChatScreenState();
}

class _ChatScreenState extends State<ChatScreen> {
  final TextEditingController _controller = TextEditingController();
  String _response = '';
  bool _loading = false;
  String? _error;

  void _send() {
    final prompt = _controller.text.trim();
    if (prompt.isEmpty || _loading) return;
    setState(() {
      _loading = true;
      _error = null;
      _response = '';
    });
    _controller.clear();
    widget.client.streamChat(prompt).listen((chunk) {
      setState(() => _response += chunk);
    }, onDone: () {
      if (mounted) setState(() => _loading = false);
    }, onError: (Object e) {
      if (mounted) {
        setState(() {
          _loading = false;
          _error = 'Request failed: $e';
        });
      }
    });
  }

  @override
  void dispose() {
    _controller.dispose();
    super.dispose();
  }

  @override
  Widget build(BuildContext context) {
    return Scaffold(
      appBar: AppBar(title: const Text('AI Chat')),
      body: Column(
        children: [
          Expanded(
            child: ListView(
              children: [
                if (_loading && _response.isEmpty && _error == null)
                  const Padding(
                    padding: EdgeInsets.all(16),
                    child: CircularProgressIndicator(key: Key('loading')),
                  ),
                if (_response.isNotEmpty)
                  Padding(
                    padding: const EdgeInsets.all(16),
                    child: Text(_response, key: const Key('response')),
                  ),
                if (_error != null)
                  Padding(
                    padding: const EdgeInsets.all(16),
                    child: Text(_error!, key: const Key('error')),
                  ),
              ],
            ),
          ),
          Padding(
            padding: const EdgeInsets.all(16),
            child: Row(
              children: [
                Expanded(
                  child: TextField(
                    key: const Key('prompt'),
                    controller: _controller,
                    decoration: const InputDecoration(hintText: 'Ask anything'),
                  ),
                ),
                FilledButton(
                  key: const Key('send'),
                  onPressed: _loading ? null : _send,
                  child: const Text('Send'),
                ),
              ],
            ),
          ),
        ],
      ),
    );
  }
}

Three visual states fall out of three fields: a spinner while _loading with no output yet, a Text that grows as chunks land, and an error message. The Keys are there for the finders you are about to write.

Widget test fundamentals

The flutter_test package gives you four pieces, and they map one-to-one onto how a human uses the app:

• testWidgets() replaces test() and hands you a fresh WidgetTester per case
• WidgetTester builds and drives widgets in a fake environment
• Finder objects locate widgets in the tree
• Matcher constants (findsOneWidget, findsNothing, findsNWidgets) verify what the finder found

The canonical skeleton:

testWidgets('renders the send button', (tester) async {
  await tester.pumpWidget(
    MaterialApp(home: ChatScreen(client: FakeLlmClient.empty())),
  );
  expect(find.byKey(const Key('send')), findsOneWidget);
});

pumpWidget builds the widget tree. After that, nothing renders on its own: when you tap a button and the widget calls setState, you must explicitly ask for a frame. That is what the pump methods are for:

• tester.pump() schedules a single frame
• tester.pump(Duration) advances the fake clock and then renders a frame, which is exactly how you step a streaming response forward
• tester.pumpAndSettle() keeps pumping until no frames are scheduled, useful for animations, but dangerous around live streams

For finders, reach for the find constant: find.text('Send'), find.byKey(const Key('loading')), find.byType(TextField), or find.byIcon(Icons.send). If the same text appears twice, find.text(...).last disambiguates.

Mocking the LLM service: a fake first

A fake is a real class that implements LlmClient and returns canned data. It is fast, readable, and requires no code generation. For streaming you want to control when each chunk is delivered, so build one that yields on a schedule:

class FakeLlmClient implements LlmClient {
  FakeLlmClient({
    required this.chunks,
    this.error,
    this.delay = Duration.zero,
  });

  final List<String> chunks;
  final Object? error;
  final Duration delay;

  @override
  Stream<String> streamChat(String prompt) async* {
    if (delay > Duration.zero) {
      await Future<void>.delayed(delay);
    }
    if (error != null) {
      throw error!;
    }
    for (final chunk in chunks) {
      yield chunk;
    }
  }
}

Because widget tests run under fake async, the Future.delayed is not real time: tester.pump(const Duration(milliseconds: 100)) advances it deterministically. If you need to keep the stream open indefinitely (to hold a loading state), back the fake with a StreamController instead:

class HangLlmClient implements LlmClient {
  final _controller = StreamController<String>();

  @override
  Stream<String> streamChat(String prompt) => _controller.stream;
}

Reach for Mockito when you need call verification or argument capture. Declare the mock and generate it:

@GenerateMocks([LlmClient])
void main() {}

Run dart run build_runner build, then stub behavior with when and assert calls with verify:

final client = MockLlmClient();
when(client.streamChat('Hi')).thenAnswer(
  (_) => Stream.fromIterable(['Hello', ' world']),
);

await tester.tap(find.byKey(const Key('send')));
await tester.pump();

verify(client.streamChat('Hi')).called(1);

mocktail offers the same API without code generation. Both are fine; the fake keeps your widget tests free of extra tooling.

Testing the loading state

The loading state is the easiest to get wrong because pumpAndSettle will hang on it: a CircularProgressIndicator animates forever, so settling never completes and the test fails with a timeout. Pump a fixed number of frames instead.

testWidgets('shows a spinner while the model is thinking', (tester) async {
  final client = HangLlmClient();
  await tester.pumpWidget(MaterialApp(home: ChatScreen(client: client)));

  await tester.enterText(find.byKey(const Key('prompt')), 'Summarize this');
  await tester.tap(find.byKey(const Key('send')));
  await tester.pump();

  expect(find.byKey(const Key('loading')), findsOneWidget);
  expect(
    tester.widget<FilledButton>(find.byKey(const Key('send'))).onPressed,
    isNull,
  );
});

Because HangLlmClient never closes its stream, the spinner must remain visible and the send button must be disabled (its onPressed is null) for the whole duration. This is the assertion that catches double-send bugs.

Testing streaming responses chunk by chunk

Here is the payoff of fake async. Feed the fake chunks with a delay, then advance the clock and assert the partial text grows:

testWidgets('renders streamed tokens as they arrive', (tester) async {
  final client = FakeLlmClient(
    chunks: ['Dart ', 'is ', 'fun.'],
    delay: const Duration(milliseconds: 50),
  );
  await tester.pumpWidget(MaterialApp(home: ChatScreen(client: client)));

  await tester.enterText(find.byKey(const Key('prompt')), 'Tell me');
  await tester.tap(find.byKey(const Key('send')));
  await tester.pump();

  expect(find.text('Dart '), findsOneWidget);

  await tester.pump(const Duration(milliseconds: 50));
  expect(find.text('Dart is '), findsOneWidget);

  await tester.pump(const Duration(milliseconds: 50));
  expect(find.text('Dart is fun.'), findsOneWidget);
  expect(find.byKey(const Key('loading')), findsNothing);
});

Notice the rhythm: pump() after the tap renders the first chunk (the stream starts synchronously up to the first await), then each pump(Duration) delivers the next chunk. The final expectation proves the spinner is gone once the stream completes. Do not reach for pumpAndSettle here; you want to observe intermediate frames, not skip past them.

If your UI debounces input or runs a cursor blink, keep those animations deterministic or gate them behind a flag, or pumpAndSettle will either time out or make assertions on blinking text flaky.

Testing the error state

A model can fail before the first token or mid-stream. The fake throws, and the widget must swap the spinner for the error message:

testWidgets('shows an error when the model call fails', (tester) async {
  final client = FakeLlmClient(
    chunks: const [],
    error: Exception('rate limited'),
  );
  await tester.pumpWidget(MaterialApp(home: ChatScreen(client: client)));

  await tester.enterText(find.byKey(const Key('prompt')), 'Hi');
  await tester.tap(find.byKey(const Key('send')));
  await tester.pump();
  await tester.pump();

  expect(find.byKey(const Key('error')), findsOneWidget);
  expect(find.textContaining('Request failed'), findsOneWidget);
  expect(find.byKey(const Key('loading')), findsNothing);
  expect(find.text('Hi'), findsNothing);
});

The first pump() schedules the frame after the tap; the second lets the stream’s error propagate into setState. textContaining is a useful matcher when you do not want to couple the test to the full message. If your app surfaces unexpected errors to the user, also assert the button is re-enabled so the user can retry.

Golden tests for AI UI

Golden tests compare a rendered widget against a committed baseline PNG, catching accidental layout drift in your chat bubbles, typing indicator, and message list. Generate the baseline once:

testWidgets('chat screen renders a completed response', (tester) async {
  final client = FakeLlmClient(chunks: ['Hello from the model.']);
  await tester.pumpWidget(MaterialApp(home: ChatScreen(client: client)));

  await tester.enterText(find.byKey(const Key('prompt')), 'Hi');
  await tester.tap(find.byKey(const Key('send')));
  await tester.pumpAndSettle();

  await expectLater(
    find.byType(ChatScreen),
    matchesGoldenFile('goldens/chat_screen.png'),
  );
});

Manage the baseline with flags:

flutter test --update-goldens
flutter test

Two rules keep goldens from being flaky. First, never capture while a stream is mid-flight or a progress indicator is animating; settle to a completed response first. Second, pin fonts. AI chat UIs frequently use a monospaced or rounded font, and text rendering differs across platforms, so set a fixed fontFamily in the ThemeData passed to the test widget, or the same source will render differently on another developer’s machine.

Putting It All Together

Here is the complete suite so you can copy, adapt, and run it:

import 'package:flutter/material.dart';
import 'package:flutter_test/flutter_test.dart';
import 'package:mockito/annotations.dart';
import 'package:mockito/mockito.dart';

import 'chat_screen.dart';
import 'chat_screen_test.mocks.dart';

@GenerateMocks([LlmClient])
void main() {
  group('ChatScreen', () {
    testWidgets('shows a spinner while the model is thinking', (tester) async {
      final client = HangLlmClient();
      await tester.pumpWidget(MaterialApp(home: ChatScreen(client: client)));

      await tester.enterText(find.byKey(const Key('prompt')), 'Summarize');
      await tester.tap(find.byKey(const Key('send')));
      await tester.pump();

      expect(find.byKey(const Key('loading')), findsOneWidget);
      expect(
        tester.widget<FilledButton>(find.byKey(const Key('send'))).onPressed,
        isNull,
      );
    });

    testWidgets('renders streamed tokens as they arrive', (tester) async {
      final client = FakeLlmClient(
        chunks: ['Dart ', 'is ', 'fun.'],
        delay: const Duration(milliseconds: 50),
      );
      await tester.pumpWidget(MaterialApp(home: ChatScreen(client: client)));

      await tester.enterText(find.byKey(const Key('prompt')), 'Tell me');
      await tester.tap(find.byKey(const Key('send')));
      await tester.pump();

      expect(find.text('Dart '), findsOneWidget);

      await tester.pump(const Duration(milliseconds: 50));
      expect(find.text('Dart is '), findsOneWidget);

      await tester.pump(const Duration(milliseconds: 50));
      expect(find.text('Dart is fun.'), findsOneWidget);
      expect(find.byKey(const Key('loading')), findsNothing);
    });

    testWidgets('shows an error when the model call fails', (tester) async {
      final client = FakeLlmClient(
        chunks: const [],
        error: Exception('rate limited'),
      );
      await tester.pumpWidget(MaterialApp(home: ChatScreen(client: client)));

      await tester.enterText(find.byKey(const Key('prompt')), 'Hi');
      await tester.tap(find.byKey(const Key('send')));
      await tester.pump();
      await tester.pump();

      expect(find.textContaining('Request failed'), findsOneWidget);
      expect(find.byKey(const Key('loading')), findsNothing);
    });

    testWidgets('forwards the prompt to the client exactly once',
        (tester) async {
      final client = MockLlmClient();
      when(client.streamChat('Hi')).thenAnswer(
        (_) => Stream.fromIterable(['Hello', ' world']),
      );
      await tester.pumpWidget(MaterialApp(home: ChatScreen(client: client)));

      await tester.enterText(find.byKey(const Key('prompt')), 'Hi');
      await tester.tap(find.byKey(const Key('send')));
      await tester.pump();

      verify(client.streamChat('Hi')).called(1);
    });
  });
}

Run the suite with flutter test and get deterministic, offline coverage of every state your model integration can enter.

Conclusion & Next Steps

Widget tests give AI features something live models cannot: a deterministic harness. With a fake or mock LlmClient, testWidgets and WidgetTester let you prove the spinner shows, tokens render incrementally, failures surface gracefully, and the UI never double-sends. The same pump loop that drives streaming also keeps golden baselines honest.

From here, level up by testing token cancellation (close the stream subscription and assert the UI stops), running the same widget tests against mocktail to drop code generation, adding integration tests with integration_test that hit a real-but-sandboxed endpoint, and extracting the state machine into a ChangeNotifier you can unit test independently of widgets. The pattern is always the same: inject the model behind an interface, fake the stream, and pump the frames.