diff --git a/dev/integration_tests/flutter_gallery/lib/demo/material/dialog_demo.dart b/dev/integration_tests/flutter_gallery/lib/demo/material/dialog_demo.dart
index 764670166c..ee0185f087 100644
--- a/dev/integration_tests/flutter_gallery/lib/demo/material/dialog_demo.dart
+++ b/dev/integration_tests/flutter_gallery/lib/demo/material/dialog_demo.dart
@@ -179,7 +179,7 @@ class DialogDemoState extends State<DialogDemo> {
                 initialTime: _selectedTime!,
               )
               .then((TimeOfDay? value) {
-                if (!mounted) {
+                if (!context.mounted) {
                   return;
                 }
                 if (value != null && value != _selectedTime) {