diff --git a/test/controls/yaru_back_button_test.dart b/test/controls/yaru_back_button_test.dart
index 1fca9d0a7..94c150156 100644
--- a/test/controls/yaru_back_button_test.dart
+++ b/test/controls/yaru_back_button_test.dart
@@ -36,7 +36,7 @@ void main() {
 }
 
 final goldenVariant = ValueVariant({
-  ...goldenThemeVariants('normal'),
-  ...goldenThemeVariants('hovered', {MaterialState.hovered: true}),
-  ...goldenThemeVariants('pressed', {MaterialState.pressed: true}),
+  ...goldenThemeVariants('normal', <MaterialState>{}),
+  ...goldenThemeVariants('hovered', {MaterialState.hovered}),
+  ...goldenThemeVariants('pressed', {MaterialState.pressed}),
 });
diff --git a/test/controls/yaru_circular_progress_indicator_test.dart b/test/controls/yaru_circular_progress_indicator_test.dart
index 48da9e84c..a91e35792 100644
--- a/test/controls/yaru_circular_progress_indicator_test.dart
+++ b/test/controls/yaru_circular_progress_indicator_test.dart
@@ -64,8 +64,8 @@ void main() {
 }
 
 final goldenVariant = ValueVariant({
-  ...goldenThemeVariants('indeterminate'),
-  ...goldenThemeVariants('empty', {}, 0.0),
-  ...goldenThemeVariants('half', {}, 0.5),
-  ...goldenThemeVariants('full', {}, 1.0),
+  ...goldenThemeVariants('indeterminate', null),
+  ...goldenThemeVariants('empty', 0.0),
+  ...goldenThemeVariants('half', 0.5),
+  ...goldenThemeVariants('full', 1.0),
 });
diff --git a/test/controls/yaru_close_button_test.dart b/test/controls/yaru_close_button_test.dart
index cc0dc2c23..a928e2c01 100644
--- a/test/controls/yaru_close_button_test.dart
+++ b/test/controls/yaru_close_button_test.dart
@@ -36,7 +36,7 @@ void main() {
 }
 
 final goldenVariant = ValueVariant({
-  ...goldenThemeVariants('normal'),
-  ...goldenThemeVariants('hovered', {MaterialState.hovered: true}),
-  ...goldenThemeVariants('pressed', {MaterialState.pressed: true}),
+  ...goldenThemeVariants('normal', <MaterialState>{}),
+  ...goldenThemeVariants('hovered', {MaterialState.hovered}),
+  ...goldenThemeVariants('pressed', {MaterialState.pressed}),
 });
diff --git a/test/controls/yaru_icon_button_test.dart b/test/controls/yaru_icon_button_test.dart
index 4c3225087..2d9b1e0bc 100644
--- a/test/controls/yaru_icon_button_test.dart
+++ b/test/controls/yaru_icon_button_test.dart
@@ -18,7 +18,7 @@ void main() {
       await tester.pumpScaffold(
         YaruIconButton(
           autofocus: variant.hasState(MaterialState.focused),
-          isSelected: variant.states[MaterialState.selected],
+          isSelected: variant.value?[MaterialState.selected],
           onPressed: variant.hasState(MaterialState.disabled) ? null : () {},
           icon: const Icon(YaruIcons.star_filled),
         ),
@@ -47,7 +47,7 @@ void main() {
 
 final goldenVariant = ValueVariant({
   // normal (non-toggle) button
-  ...goldenThemeVariants('normal'),
+  ...goldenThemeVariants('normal', <MaterialState, bool>{}),
   ...goldenThemeVariants('disabled', {MaterialState.disabled: true}),
   ...goldenThemeVariants('focused', {MaterialState.focused: true}),
   ...goldenThemeVariants('hovered', {MaterialState.hovered: true}),
diff --git a/test/controls/yaru_linear_progress_indicator_test.dart b/test/controls/yaru_linear_progress_indicator_test.dart
index dec444edc..8895d7779 100644
--- a/test/controls/yaru_linear_progress_indicator_test.dart
+++ b/test/controls/yaru_linear_progress_indicator_test.dart
@@ -63,8 +63,8 @@ void main() {
 }
 
 final goldenVariant = ValueVariant({
-  ...goldenThemeVariants('indeterminate'),
-  ...goldenThemeVariants('empty', {}, 0.0),
-  ...goldenThemeVariants('half', {}, 0.5),
-  ...goldenThemeVariants('full', {}, 1.0),
+  ...goldenThemeVariants('indeterminate', null),
+  ...goldenThemeVariants('empty', 0.0),
+  ...goldenThemeVariants('half', 0.5),
+  ...goldenThemeVariants('full', 1.0),
 });
diff --git a/test/controls/yaru_option_button_test.dart b/test/controls/yaru_option_button_test.dart
index f6ac4a043..8ba0b1393 100644
--- a/test/controls/yaru_option_button_test.dart
+++ b/test/controls/yaru_option_button_test.dart
@@ -86,9 +86,9 @@ void main() {
 }
 
 final goldenVariant = ValueVariant({
-  ...goldenThemeVariants('normal'),
-  ...goldenThemeVariants('disabled', {MaterialState.disabled: true}),
-  ...goldenThemeVariants('focused', {MaterialState.focused: true}),
-  ...goldenThemeVariants('hovered', {MaterialState.hovered: true}),
-  ...goldenThemeVariants('pressed', {MaterialState.pressed: true}),
+  ...goldenThemeVariants('normal', <MaterialState>{}),
+  ...goldenThemeVariants('disabled', {MaterialState.disabled}),
+  ...goldenThemeVariants('focused', {MaterialState.focused}),
+  ...goldenThemeVariants('hovered', {MaterialState.hovered}),
+  ...goldenThemeVariants('pressed', {MaterialState.pressed}),
 });
diff --git a/test/yaru_golden_tester.dart b/test/yaru_golden_tester.dart
index 5c2af9c26..1d0554518 100644
--- a/test/yaru_golden_tester.dart
+++ b/test/yaru_golden_tester.dart
@@ -42,59 +42,56 @@ extension YaruGoldenTester on WidgetTester {
 }
 
 @immutable
-class YaruGoldenVariant {
+class YaruGoldenVariant<T> {
   YaruGoldenVariant({
     required String label,
     required this.themeMode,
-    this.states = const {},
     this.value,
   }) : label = '$label-${themeMode.name}';
 
   final String label;
   final ThemeMode themeMode;
-  final Map<MaterialState, bool> states;
-  final dynamic value;
-
-  bool hasState(MaterialState state) => states[state] == true;
+  final T? value;
 
   @override
   bool operator ==(Object other) {
     if (identical(this, other)) return true;
-    final mapEquals = const MapEquality().equals;
+    final deepEquals = const DeepCollectionEquality().equals;
     return other is YaruGoldenVariant &&
         other.label == label &&
         other.themeMode == themeMode &&
-        mapEquals(other.states, states) &&
-        other.value == value;
+        deepEquals(other.value, value);
   }
 
   @override
   int get hashCode {
-    final mapHash = const MapEquality().hash;
-    return Object.hash(label, themeMode, mapHash(states), value);
+    final deepHash = const DeepCollectionEquality().hash;
+    return Object.hash(label, themeMode, deepHash(value));
   }
 
   @override
-  String toString() =>
-      '$label: themeMode: $themeMode, states: $states, value: $value';
+  String toString() => '$label: themeMode: $themeMode, value: $value';
+}
+
+extension YaruGoldenVariantStateSet on YaruGoldenVariant<Set<MaterialState>> {
+  bool hasState(MaterialState state) => value?.contains(state) == true;
+}
+
+extension YaruGoldenVariantStateMap
+    on YaruGoldenVariant<Map<MaterialState, bool>> {
+  bool hasState(MaterialState state) => value?[state] == true;
 }
 
-List<YaruGoldenVariant> goldenThemeVariants(
-  String label, [
-  Map<MaterialState, bool> states = const {},
-  dynamic value,
-]) {
+List<YaruGoldenVariant<T>> goldenThemeVariants<T>(String label, [T? value]) {
   return [
     YaruGoldenVariant(
       label: label,
       themeMode: ThemeMode.light,
-      states: states,
       value: value,
     ),
     YaruGoldenVariant(
       label: label,
       themeMode: ThemeMode.dark,
-      states: states,
       value: value,
     ),
   ];