[Story of Your Layout] Make dropdown align to their axis and make flex widget clickable (#1392)

* Make dropdown align to their axis and make flex widget clickable

* Fix relative import and dartfmt

* Update golden files (might file because running from mac)

* Avoid changing textDirection and document onChange type signature

* Update goldens

* Change inheritFromWidgetOfExactType to dependOnInheritedWidgetOfExactType due to deprecation
diff --git a/packages/devtools_app/lib/src/flutter/controllers.dart b/packages/devtools_app/lib/src/flutter/controllers.dart
index cb4905c..afd0f2c 100644
--- a/packages/devtools_app/lib/src/flutter/controllers.dart
+++ b/packages/devtools_app/lib/src/flutter/controllers.dart
@@ -67,7 +67,7 @@
 
   static ProvidedControllers of(BuildContext context) {
     final _InheritedProvider inherited =
-        context.inheritFromWidgetOfExactType(_InheritedProvider);
+        context.dependOnInheritedWidgetOfExactType(aspect: _InheritedProvider);
     return inherited.data;
   }
 }
diff --git a/packages/devtools_app/lib/src/inspector/flutter/inspector_data_models.dart b/packages/devtools_app/lib/src/inspector/flutter/inspector_data_models.dart
index ac1b994..65de1b5 100644
--- a/packages/devtools_app/lib/src/inspector/flutter/inspector_data_models.dart
+++ b/packages/devtools_app/lib/src/inspector/flutter/inspector_data_models.dart
@@ -10,7 +10,6 @@
 import '../../utils.dart';
 import '../diagnostics_node.dart';
 import '../enum_utils.dart';
-import '../inspector_tree.dart';
 import 'story_of_your_layout/utils.dart';
 
 const Type boxConstraintsType = BoxConstraints;
@@ -95,19 +94,19 @@
 // TODO(albertusangga): Move this to [RemoteDiagnosticsNode] once dart:html app is removed
 class LayoutProperties {
   LayoutProperties(this.node, {int copyLevel = 1})
-      : description = node.diagnostic?.description,
-        size = deserializeSize(node.diagnostic?.size),
-        constraints = deserializeConstraints(node.diagnostic?.constraints),
-        isFlex = node.diagnostic?.isFlex,
-        flexFactor = node.diagnostic?.flexFactor,
+      : description = node?.description,
+        size = deserializeSize(node?.size),
+        constraints = deserializeConstraints(node?.constraints),
+        isFlex = node?.isFlex,
+        flexFactor = node?.flexFactor,
         children = copyLevel == 0
             ? []
-            : node.children
+            : node?.childrenNow
                 ?.map((child) =>
                     LayoutProperties(child, copyLevel: copyLevel - 1))
                 ?.toList(growable: false);
 
-  final InspectorTreeNode node;
+  final RemoteDiagnosticsNode node;
   final List<LayoutProperties> children;
   final BoxConstraints constraints;
   final String description;
@@ -178,7 +177,7 @@
 /// TODO(albertusangga): Move this to [RemoteDiagnosticsNode] once dart:html app is removed
 class FlexLayoutProperties extends LayoutProperties {
   FlexLayoutProperties._(
-    InspectorTreeNode node, {
+    RemoteDiagnosticsNode node, {
     this.direction,
     this.mainAxisAlignment,
     this.mainAxisSize,
@@ -188,16 +187,15 @@
     this.textBaseline,
   }) : super(node);
 
-  factory FlexLayoutProperties.fromNode(InspectorTreeNode node) {
+  factory FlexLayoutProperties.fromDiagnostics(RemoteDiagnosticsNode node) {
     // Cache the properties on an expando so that local tweaks to
     // FlexLayoutProperties persist across multiple lookups from an
-    // InspectorTreeNode.
+    // RemoteDiagnosticsNode.
     return _flexLayoutExpando[node] ??= _buildNode(node);
   }
 
-  static FlexLayoutProperties _buildNode(InspectorTreeNode node) {
-    final Map<String, Object> renderObjectJson =
-        node.diagnostic.json['renderObject'];
+  static FlexLayoutProperties _buildNode(RemoteDiagnosticsNode node) {
+    final Map<String, Object> renderObjectJson = node.json['renderObject'];
     final List<dynamic> properties = renderObjectJson['properties'];
     final Map<String, Object> data = Map<String, Object>.fromIterable(
       properties,
diff --git a/packages/devtools_app/lib/src/inspector/flutter/inspector_screen_details_tab.dart b/packages/devtools_app/lib/src/inspector/flutter/inspector_screen_details_tab.dart
index 391762e..2675234 100644
--- a/packages/devtools_app/lib/src/inspector/flutter/inspector_screen_details_tab.dart
+++ b/packages/devtools_app/lib/src/inspector/flutter/inspector_screen_details_tab.dart
@@ -6,8 +6,8 @@
 import 'package:flutter/rendering.dart';
 import 'package:flutter/widgets.dart';
 
+import '../diagnostics_node.dart';
 import '../inspector_controller.dart';
-import '../inspector_tree.dart';
 import 'inspector_data_models.dart';
 import 'story_of_your_layout/flex.dart';
 
@@ -101,7 +101,7 @@
     with AutomaticKeepAliveClientMixin<LayoutDetailsTab> {
   InspectorController get controller => widget.controller;
 
-  InspectorTreeNode get selected => controller?.selectedNode;
+  RemoteDiagnosticsNode get selected => controller?.selectedNode?.diagnostic;
 
   void onSelectionChanged() {
     setState(() {});
@@ -122,16 +122,15 @@
   @override
   Widget build(BuildContext context) {
     super.build(context);
-    final diagnostic = selected?.diagnostic;
     // TODO(albertusangga): Visualize non-flex widget constraint model
-    if (diagnostic == null ||
-        (!diagnostic.isFlex && !(diagnostic.parent?.isFlex ?? false)))
+    if (selected == null ||
+        (!selected.isFlex && !(selected.parent?.isFlex ?? false)))
       return const SizedBox();
-    final flexLayoutProperties = FlexLayoutProperties.fromNode(
-        diagnostic.isFlex ? selected : selected.parent);
-    final highlightChild = diagnostic.isFlex
-        ? null
-        : diagnostic.parent.childrenNow.indexOf(diagnostic);
+    final flexLayoutProperties = FlexLayoutProperties.fromDiagnostics(
+      selected.isFlex ? selected : selected.parent,
+    );
+    final highlightChild =
+        selected.isFlex ? null : selected.parent.childrenNow.indexOf(selected);
     return StoryOfYourFlexWidget(
       // TODO(albertusangga): Cache this instead of recomputing every build,
       flexLayoutProperties,
diff --git a/packages/devtools_app/lib/src/inspector/flutter/inspector_tree_flutter.dart b/packages/devtools_app/lib/src/inspector/flutter/inspector_tree_flutter.dart
index 7eb8ec5..6e3ee8d 100644
--- a/packages/devtools_app/lib/src/inspector/flutter/inspector_tree_flutter.dart
+++ b/packages/devtools_app/lib/src/inspector/flutter/inspector_tree_flutter.dart
@@ -569,7 +569,7 @@
                   constraintDisplayController != null)
                 ConstraintsDescription(
                   listenable: constraintDisplayController,
-                  properties: LayoutProperties(node),
+                  properties: LayoutProperties(node.diagnostic),
                 ),
             ],
           ),
diff --git a/packages/devtools_app/lib/src/inspector/flutter/story_of_your_layout/flex.dart b/packages/devtools_app/lib/src/inspector/flutter/story_of_your_layout/flex.dart
index cd1da7a..46fb672 100644
--- a/packages/devtools_app/lib/src/inspector/flutter/story_of_your_layout/flex.dart
+++ b/packages/devtools_app/lib/src/inspector/flutter/story_of_your_layout/flex.dart
@@ -35,8 +35,8 @@
 /// Hardcoded arrow size respective to its cross axis (because it's unconstrained).
 const heightAndConstraintIndicatorSize = 48.0;
 const widthAndConstraintIndicatorSize = 42.0;
-const mainAxisArrowIndicatorSize = 32.0;
-const crossAxisArrowIndicatorSize = 32.0;
+const mainAxisArrowIndicatorSize = 48.0;
+const crossAxisArrowIndicatorSize = 48.0;
 
 const heightOnlyIndicatorSize = 24.0;
 const widthOnlyIndicatorSize = 24.0;
@@ -44,8 +44,35 @@
 const largeTextScaleFactor = 1.2;
 const smallTextScaleFactor = 0.8;
 
+/// height for limiting asset image (selected one in the drop down).
 const axisAlignmentAssetImageHeight = 24.0;
-const dropdownMaxWidth = 320.0;
+
+/// width for limiting asset image (when drop down menu is open for the vertical).
+const axisAlignmentAssetImageWidth = 96.0;
+const dropdownMaxSize = 300.0;
+
+Color activeBackgroundColor(ThemeData theme) => theme.backgroundColor;
+
+Color inActiveBackgroundColor(ThemeData theme) => theme.cardColor;
+
+// Story of Layout colors
+const mainAxisLightColor = Color(0xFFF597A8);
+const mainAxisDarkColor = Color(0xFFEA637C);
+const mainAxisColor = ThemedColor(mainAxisLightColor, mainAxisDarkColor);
+
+const crossAxisLightColor = Color(0xFFB3D25A);
+const crossAxisDarkColor = Color(0xFFB3D25A);
+const crossAxisColor = ThemedColor(crossAxisLightColor, crossAxisDarkColor);
+
+const mainAxisLightTextColor = Color(0xFF913549);
+const mainAxisDarkTextColor = Color(0xFFEA637C);
+const mainAxisTextColor =
+    ThemedColor(mainAxisLightTextColor, mainAxisDarkTextColor);
+
+const crossAxisLightTextColor = Color(0xFF66672C);
+const crossAxisDarkTextColor = Color(0xFFB3D25A);
+const crossAxisTextColor =
+    ThemedColor(crossAxisLightTextColor, crossAxisDarkTextColor);
 
 class StoryOfYourFlexWidget extends StatefulWidget {
   const StoryOfYourFlexWidget(
@@ -86,6 +113,12 @@
   Color get verticalColor =>
       properties.isMainAxisVertical ? mainAxisColor : crossAxisColor;
 
+  Color get horizontalTextColor =>
+      properties.isMainAxisHorizontal ? mainAxisTextColor : crossAxisTextColor;
+
+  Color get verticalTextColor =>
+      properties.isMainAxisVertical ? mainAxisTextColor : crossAxisTextColor;
+
   String get flexType => properties.type;
 
   MainAxisAlignment get mainAxisAlignment => properties.mainAxisAlignment;
@@ -157,6 +190,17 @@
     );
   }
 
+  Function _onTap(LayoutProperties properties) => () async {
+        final controller = widget.inspectorController;
+        final diagnostic = properties.node;
+        controller.refreshSelection(diagnostic, diagnostic, false);
+        final inspectorService = await diagnostic.inspectorService;
+        await inspectorService.setSelectionInspector(
+          diagnostic.valueRef,
+          true,
+        );
+      };
+
   Widget _visualizeChild({
     LayoutProperties childProperties,
     Color backgroundColor,
@@ -170,16 +214,7 @@
       top: renderOffset.dy,
       left: renderOffset.dx,
       child: InkWell(
-        onTap: () async {
-          final controller = widget.inspectorController;
-          final diagnostic = childProperties.node.diagnostic;
-          // TODO(albertusangga) fix/investigate why calling setSelectedNode is not sufficient
-          controller.refreshSelection(diagnostic, diagnostic, false);
-          controller.setSelectedNode(childProperties.node);
-          final inspectorService = await diagnostic.inspectorService;
-          await inspectorService.setSelectionInspector(
-              diagnostic.valueRef, true);
-        },
+        onTap: _onTap(childProperties),
         child: Container(
           width: renderSize.width,
           height: renderSize.height,
@@ -259,8 +294,8 @@
               _visualizeChild(
                 backgroundColor:
                     widget.highlightChild != null && i == widget.highlightChild
-                        ? theme.backgroundColor
-                        : theme.cardColor,
+                        ? activeBackgroundColor(theme)
+                        : inActiveBackgroundColor(theme),
                 childProperties: children[i],
                 borderColor: i.isOdd ? mainAxisColor : crossAxisColor,
                 textColor: i.isOdd ? null : const Color(0xFF303030),
@@ -318,17 +353,9 @@
   }
 
   Widget _buildAxisAlignmentDropdown(Axis axis) {
-    Color color;
-    String axisDescription;
+    final color = axis == direction ? mainAxisTextColor : crossAxisTextColor;
     List<Object> alignmentEnumEntries;
     Object selected;
-    if (axis == Axis.horizontal) {
-      color = horizontalColor;
-      axisDescription = properties.horizontalDirectionDescription;
-    } else {
-      color = verticalColor;
-      axisDescription = properties.verticalDirectionDescription;
-    }
     if (axis == direction) {
       alignmentEnumEntries = MainAxisAlignment.values;
       selected = mainAxisAlignment;
@@ -340,70 +367,103 @@
       }
       selected = crossAxisAlignment;
     }
-    return Container(
-      constraints: const BoxConstraints(maxWidth: dropdownMaxWidth),
-      child: Column(
-        mainAxisAlignment: MainAxisAlignment.start,
-        crossAxisAlignment: CrossAxisAlignment.start,
-        mainAxisSize: MainAxisSize.min,
-        children: <Widget>[
-          Text(
-            '$axisDescription Alignment: ',
-            textScaleFactor: largeTextScaleFactor,
-          ),
-          Container(
-            margin: const EdgeInsets.only(left: 8.0),
-            child: DropdownButton(
-              isExpanded: true,
-              value: selected,
-              items: [
-                for (var alignment in alignmentEnumEntries)
-                  DropdownMenuItem(
-                    value: alignment,
-                    child: Row(
-                      mainAxisAlignment: MainAxisAlignment.end,
-                      crossAxisAlignment: CrossAxisAlignment.center,
-                      children: <Widget>[
-                        Expanded(
-                          child: Container(
-                            margin: const EdgeInsets.only(right: 8.0),
-                            child: Text(
-                              describeEnum(alignment),
-                              style: TextStyle(color: color),
-                              textAlign: TextAlign.right,
-                            ),
-                          ),
+    return RotatedBox(
+      quarterTurns: axis == Axis.vertical ? 3 : 0,
+      child: Container(
+        constraints: const BoxConstraints(
+            maxWidth: dropdownMaxSize, maxHeight: dropdownMaxSize),
+        child: DropdownButton(
+          value: selected,
+          isExpanded: true,
+          itemHeight: axis == Axis.vertical ? 100.0 : null,
+          selectedItemBuilder: (context) {
+            return [
+              for (var alignment in alignmentEnumEntries)
+                Row(
+                  mainAxisAlignment: MainAxisAlignment.end,
+                  crossAxisAlignment: CrossAxisAlignment.center,
+                  children: <Widget>[
+                    Expanded(
+                      child: Container(
+                        child: Text(
+                          describeEnum(alignment),
+                          style: TextStyle(color: color),
+                          textAlign: TextAlign.center,
                         ),
-                        Image.asset(
+                      ),
+                    ),
+                    Expanded(
+                      child: RotatedBox(
+                        quarterTurns: axis == Axis.vertical ? 2 : 0,
+                        child: Image.asset(
                           (axis == direction)
                               ? mainAxisAssetImageUrl(alignment)
                               : crossAxisAssetImageUrl(alignment),
                           height: axisAlignmentAssetImageHeight,
                           fit: BoxFit.contain,
                         ),
-                      ],
+                      ),
                     ),
-                  )
-              ],
-              onChanged: (Object newSelection) async {
-                if (axis == direction) {
-                  properties.mainAxisAlignment = newSelection;
-                } else {
-                  properties.crossAxisAlignment = newSelection;
-                }
-                final service =
-                    await properties.node.diagnostic.inspectorService;
-                final arg = properties.node.diagnostic.valueRef;
-                await service.invokeTweakFlexProperties(
-                  arg,
-                  properties.mainAxisAlignment,
-                  properties.crossAxisAlignment,
-                );
-                setState(() {});
-              },
-            ),
-          )
-        ],
+                  ],
+                )
+            ];
+          },
+          items: [
+            for (var alignment in alignmentEnumEntries)
+              DropdownMenuItem(
+                value: alignment,
+                child: Container(
+                  padding: const EdgeInsets.symmetric(vertical: margin),
+                  child: Row(
+                    mainAxisAlignment: MainAxisAlignment.end,
+                    crossAxisAlignment: CrossAxisAlignment.center,
+                    children: <Widget>[
+                      Expanded(
+                        child: Container(
+                          child: Text(
+                            describeEnum(alignment),
+                            style: TextStyle(color: color),
+                            textAlign: TextAlign.center,
+                          ),
+                        ),
+                      ),
+                      Expanded(
+                        child: RotatedBox(
+                          quarterTurns: axis == Axis.vertical ? 1 : 0,
+                          child: Image.asset(
+                            (axis == direction)
+                                ? mainAxisAssetImageUrl(alignment)
+                                : crossAxisAssetImageUrl(alignment),
+                            width: axisAlignmentAssetImageWidth,
+                            fit: BoxFit.contain,
+                          ),
+                        ),
+                      ),
+                    ],
+                  ),
+                ),
+              )
+          ],
+          onChanged: (Object newSelection) async {
+            // newSelection is an object instead of type here because
+            // the type is dependent on the `axis` parameter
+            // if the axis is the main axis the type should be [MainAxisAlignment]
+            // if the axis is the cross axis the type should be [CrossAxisAlignment]
+            if (axis == direction) {
+              properties.mainAxisAlignment = newSelection;
+            } else {
+              properties.crossAxisAlignment = newSelection;
+            }
+            final service = await properties.node.inspectorService;
+            final arg = properties.node.valueRef;
+            await service.invokeTweakFlexProperties(
+              arg,
+              properties.mainAxisAlignment,
+              properties.crossAxisAlignment,
+            );
+            setState(() {});
+          },
+        ),
       ),
     );
   }
@@ -412,25 +472,20 @@
   Widget build(BuildContext context) {
     final theme = Theme.of(context);
     return Container(
-      padding: const EdgeInsets.only(top: 32.0),
       child: Column(
         mainAxisAlignment: MainAxisAlignment.start,
         children: [
           Container(
-            margin: const EdgeInsets.only(bottom: 24.0),
             child: Text(
               'Story of the flex layout of your $flexType widget',
               style: theme.textTheme.headline,
               textAlign: TextAlign.center,
             ),
           ),
-          Align(
-            alignment: Alignment.centerRight,
-            child: _buildAxisAlignmentDropdown(Axis.horizontal),
-          ),
-          Flexible(
+          Expanded(
             child: Container(
               margin: const EdgeInsets.all(margin),
+              padding: const EdgeInsets.only(bottom: margin, right: margin),
               child: LayoutBuilder(builder: (context, constraints) {
                 final maxHeight = constraints.maxHeight;
                 final maxWidth = constraints.maxWidth;
@@ -446,27 +501,33 @@
                             top: mainAxisArrowIndicatorSize,
                             left: crossAxisArrowIndicatorSize + margin,
                           ),
-                          child: WidgetVisualizer(
-                            title: flexType,
-                            hint: Container(
-                              padding: const EdgeInsets.all(4.0),
-                              child: Text(
-                                'Total Flex Factor: ${properties?.totalFlex}',
-                                textScaleFactor: largeTextScaleFactor,
-                                style: TextStyle(
-                                  fontWeight: FontWeight.bold,
+                          child: InkWell(
+                            onTap: _onTap(properties),
+                            child: WidgetVisualizer(
+                              title: flexType,
+                              backgroundColor: widget.highlightChild == null
+                                  ? activeBackgroundColor(theme)
+                                  : null,
+                              hint: Container(
+                                padding: const EdgeInsets.all(4.0),
+                                child: Text(
+                                  'Total Flex Factor: ${properties?.totalFlex}',
+                                  textScaleFactor: largeTextScaleFactor,
+                                  style: TextStyle(
+                                    fontWeight: FontWeight.bold,
+                                  ),
                                 ),
                               ),
-                            ),
-                            borderColor: mainAxisColor,
-                            child: Container(
-                              margin: const EdgeInsets.only(
-                                /// margin for the outer width/height
-                                ///  so that they don't stick to the corner
-                                right: margin,
-                                bottom: margin,
+                              borderColor: mainAxisColor,
+                              child: Container(
+                                margin: const EdgeInsets.only(
+                                  /// margin for the outer width/height
+                                  ///  so that they don't stick to the corner
+                                  right: margin,
+                                  bottom: margin,
+                                ),
+                                child: _visualizeFlex(context),
                               ),
-                              child: _visualizeFlex(context),
                             ),
                           ),
                         ),
@@ -476,17 +537,26 @@
                         child: Container(
                           height: maxHeight - mainAxisArrowIndicatorSize,
                           width: crossAxisArrowIndicatorSize,
-                          child: ArrowWrapper.unidirectional(
-                            arrowColor: verticalColor,
-                            child: RotatedBox(
-                              quarterTurns: 3,
-                              child: Text(
-                                properties.verticalDirectionDescription,
-                                textAlign: TextAlign.center,
-                                textScaleFactor: largeTextScaleFactor,
+                          child: Column(
+                            children: <Widget>[
+                              Expanded(
+                                child: ArrowWrapper.unidirectional(
+                                  arrowColor: verticalColor,
+                                  child: RotatedBox(
+                                    quarterTurns: 3,
+                                    child: Text(
+                                      properties.verticalDirectionDescription,
+                                      textAlign: TextAlign.center,
+                                      textScaleFactor: largeTextScaleFactor,
+                                      style:
+                                          TextStyle(color: verticalTextColor),
+                                    ),
+                                  ),
+                                  type: ArrowType.down,
+                                ),
                               ),
-                            ),
-                            type: ArrowType.down,
+                              _buildAxisAlignmentDropdown(Axis.vertical),
+                            ],
                           ),
                         ),
                       ),
@@ -496,16 +566,25 @@
                           height: mainAxisArrowIndicatorSize,
                           width:
                               maxWidth - crossAxisArrowIndicatorSize - margin,
-                          child: ArrowWrapper.unidirectional(
-                            arrowColor: horizontalColor,
-                            child: FittedBox(
-                              child: Text(
-                                properties.horizontalDirectionDescription,
-                                textAlign: TextAlign.center,
-                                textScaleFactor: largeTextScaleFactor,
+                          child: Row(
+                            children: <Widget>[
+                              Expanded(
+                                child: ArrowWrapper.unidirectional(
+                                  arrowColor: horizontalColor,
+                                  child: FittedBox(
+                                    child: Text(
+                                      properties.horizontalDirectionDescription,
+                                      textAlign: TextAlign.center,
+                                      textScaleFactor: largeTextScaleFactor,
+                                      style:
+                                          TextStyle(color: horizontalTextColor),
+                                    ),
+                                  ),
+                                  type: ArrowType.right,
+                                ),
                               ),
-                            ),
-                            type: ArrowType.right,
+                              _buildAxisAlignmentDropdown(Axis.horizontal),
+                            ],
                           ),
                         ),
                       ),
@@ -515,13 +594,6 @@
               }),
             ),
           ),
-          Container(
-            margin: const EdgeInsets.only(left: 8.0),
-            child: Align(
-              alignment: Alignment.centerLeft,
-              child: _buildAxisAlignmentDropdown(Axis.vertical),
-            ),
-          ),
         ],
       ),
     );
diff --git a/packages/devtools_app/lib/src/ui/colors.dart b/packages/devtools_app/lib/src/ui/colors.dart
index 47cef78..851cb59 100644
--- a/packages/devtools_app/lib/src/ui/colors.dart
+++ b/packages/devtools_app/lib/src/ui/colors.dart
@@ -96,12 +96,3 @@
   Colors.black54,
   Color.fromARGB(255, 200, 200, 200),
 );
-
-// Story of Layout colors
-const mainAxisLightColor = Color(0xFFF597A8);
-const mainAxisDarkColor = Color(0xFFEA637C);
-const mainAxisColor = ThemedColor(mainAxisLightColor, mainAxisDarkColor);
-
-const crossAxisLightColor = Color(0xFFB3D25A);
-const crossAxisDarkColor = Color(0xFFB3D25A);
-const crossAxisColor = ThemedColor(crossAxisLightColor, crossAxisDarkColor);
diff --git a/packages/devtools_app/test/flutter/inspector_screen_test.dart b/packages/devtools_app/test/flutter/inspector_screen_test.dart
index cead3a5..1c530bb 100644
--- a/packages/devtools_app/test/flutter/inspector_screen_test.dart
+++ b/packages/devtools_app/test/flutter/inspector_screen_test.dart
@@ -231,7 +231,7 @@
       await tester.pumpWidget(
         MaterialApp(
           home: ConstraintsDescription(
-            properties: LayoutProperties(node),
+            properties: LayoutProperties(node.diagnostic),
             listenable: animationController,
           ),
         ),
diff --git a/packages/devtools_app/test/flutter/story_of_layout/flex_test.dart b/packages/devtools_app/test/flutter/story_of_layout/flex_test.dart
index b1bf7cc..d0eddf0 100644
--- a/packages/devtools_app/test/flutter/story_of_layout/flex_test.dart
+++ b/packages/devtools_app/test/flutter/story_of_layout/flex_test.dart
@@ -3,29 +3,16 @@
 // found in the LICENSE file.
 
 import 'dart:convert';
-import 'dart:typed_data';
 
 import 'package:devtools_app/src/inspector/diagnostics_node.dart';
 import 'package:devtools_app/src/inspector/flutter/inspector_data_models.dart';
 import 'package:devtools_app/src/inspector/flutter/story_of_your_layout/flex.dart';
-import 'package:devtools_app/src/inspector/inspector_tree.dart';
 import 'package:flutter/foundation.dart';
 import 'package:flutter/material.dart';
-import 'package:flutter/services.dart';
 import 'package:flutter_test/flutter_test.dart';
 
 import '../wrappers.dart';
 
-// TODO(albertusangga) Investigate flaky asset test
-/// This is current workaround for golden testing asset
-/// Currently we are preventing testing the assets since it seems to be flaky.
-class MockAssetBundle extends CachingAssetBundle {
-  @override
-  Future<ByteData> load(String key) async {
-    return ByteData.view(Uint8List.fromList([0]).buffer);
-  }
-}
-
 void main() {
   const windowSize = Size(1750, 1750);
 
@@ -266,15 +253,11 @@
 
   testWidgets('Row golden test', (WidgetTester tester) async {
     final rowWidgetJsonNode = buildDiagnosticsNodeJson(Axis.horizontal);
-    final diagnostic =
+    final diagnostics =
         RemoteDiagnosticsNode(rowWidgetJsonNode, null, false, null);
-    final node = InspectorTreeNode()..diagnostic = diagnostic;
-    for (var child in diagnostic.childrenNow) {
-      node.appendChild(InspectorTreeNode()..diagnostic = child);
-    }
     await setWindowSize(windowSize);
-    final widget =
-        wrap(StoryOfYourFlexWidget(FlexLayoutProperties.fromNode(node)));
+    final widget = wrap(StoryOfYourFlexWidget(
+        FlexLayoutProperties.fromDiagnostics(diagnostics)));
     await pump(tester, widget);
     await expectLater(
       find.byWidget(widget),
@@ -284,15 +267,11 @@
 
   testWidgets('Column golden test', (WidgetTester tester) async {
     final columnWidgetJsonNode = buildDiagnosticsNodeJson(Axis.vertical);
-    final diagnostic =
+    final diagnostics =
         RemoteDiagnosticsNode(columnWidgetJsonNode, null, false, null);
-    final node = InspectorTreeNode()..diagnostic = diagnostic;
-    for (var child in diagnostic.childrenNow) {
-      node.appendChild(InspectorTreeNode()..diagnostic = child);
-    }
     await setWindowSize(windowSize);
-    final widget =
-        wrap(StoryOfYourFlexWidget(FlexLayoutProperties.fromNode(node)));
+    final widget = wrap(StoryOfYourFlexWidget(
+        FlexLayoutProperties.fromDiagnostics(diagnostics)));
     await pump(tester, widget);
     await expectLater(
       find.byWidget(widget),
diff --git a/packages/devtools_app/test/flutter/story_of_layout/goldens/story_of_column_layout.png b/packages/devtools_app/test/flutter/story_of_layout/goldens/story_of_column_layout.png
index e3bbe8d..a7295c0 100644
--- a/packages/devtools_app/test/flutter/story_of_layout/goldens/story_of_column_layout.png
+++ b/packages/devtools_app/test/flutter/story_of_layout/goldens/story_of_column_layout.png
Binary files differ
diff --git a/packages/devtools_app/test/flutter/story_of_layout/goldens/story_of_row_layout.png b/packages/devtools_app/test/flutter/story_of_layout/goldens/story_of_row_layout.png
index 144b144..fcb9763 100644
--- a/packages/devtools_app/test/flutter/story_of_layout/goldens/story_of_row_layout.png
+++ b/packages/devtools_app/test/flutter/story_of_layout/goldens/story_of_row_layout.png
Binary files differ
diff --git a/packages/devtools_app/test/flutter/story_of_layout/inspector_data_models_test.dart b/packages/devtools_app/test/flutter/story_of_layout/inspector_data_models_test.dart
index b088614..d50e962 100644
--- a/packages/devtools_app/test/flutter/story_of_layout/inspector_data_models_test.dart
+++ b/packages/devtools_app/test/flutter/story_of_layout/inspector_data_models_test.dart
@@ -7,7 +7,6 @@
 import 'package:devtools_app/src/inspector/diagnostics_node.dart';
 import 'package:devtools_app/src/inspector/flutter/inspector_data_models.dart';
 import 'package:devtools_app/src/inspector/flutter/story_of_your_layout/utils.dart';
-import 'package:devtools_app/src/inspector/inspector_tree.dart';
 import 'package:flutter/widgets.dart';
 import 'package:test/test.dart';
 
@@ -172,11 +171,10 @@
         ]
       }
     ''');
-    final diagnostic =
+    final diagnostics =
         RemoteDiagnosticsNode({'renderObject': flexJson}, null, null, null);
-    final node = InspectorTreeNode()..diagnostic = diagnostic;
     final FlexLayoutProperties flexProperties =
-        FlexLayoutProperties.fromNode(node);
+        FlexLayoutProperties.fromDiagnostics(diagnostics);
     expect(flexProperties.direction, Axis.horizontal);
     expect(flexProperties.mainAxisAlignment, MainAxisAlignment.start);
     expect(flexProperties.mainAxisSize, MainAxisSize.max);
@@ -302,9 +300,8 @@
         "widgetRuntimeType": "Row"
     }
     ''');
-      final diagnostic = RemoteDiagnosticsNode(json, null, false, null);
-      final layoutProperties =
-          LayoutProperties(InspectorTreeNode()..diagnostic = diagnostic);
+      final diagnostics = RemoteDiagnosticsNode(json, null, false, null);
+      final layoutProperties = LayoutProperties(diagnostics);
 
       expect(layoutProperties.size, const Size(432.0, 56.0));
       expect(
@@ -334,8 +331,8 @@
               }
             }
           ''');
-        final layoutProperties = LayoutProperties(InspectorTreeNode()
-          ..diagnostic = RemoteDiagnosticsNode(json, null, false, null));
+        final layoutProperties =
+            LayoutProperties(RemoteDiagnosticsNode(json, null, false, null));
         expect(layoutProperties.describeHeightConstraints(), 'h=56.0');
         expect(layoutProperties.describeWidthConstraints(), 'w=25.0');
       });
@@ -355,8 +352,8 @@
               }
             }
           ''');
-        final layoutProperties = LayoutProperties(InspectorTreeNode()
-          ..diagnostic = RemoteDiagnosticsNode(json, null, false, null));
+        final layoutProperties =
+            LayoutProperties(RemoteDiagnosticsNode(json, null, false, null));
         expect(layoutProperties.describeHeightConstraints(), '75.0<=h<=100.0');
         expect(layoutProperties.describeWidthConstraints(), '25.0<=w<=50.0');
       });
@@ -374,8 +371,8 @@
               }
             }
           ''');
-        final layoutProperties = LayoutProperties(InspectorTreeNode()
-          ..diagnostic = RemoteDiagnosticsNode(json, null, false, null));
+        final layoutProperties =
+            LayoutProperties(RemoteDiagnosticsNode(json, null, false, null));
         expect(layoutProperties.describeHeightConstraints(), 'h=unconstrained');
         expect(layoutProperties.describeWidthConstraints(), 'w=unconstrained');
       });
@@ -392,8 +389,8 @@
               }
             }
           ''');
-      final layoutProperties = LayoutProperties(InspectorTreeNode()
-        ..diagnostic = RemoteDiagnosticsNode(json, null, false, null));
+      final layoutProperties =
+          LayoutProperties(RemoteDiagnosticsNode(json, null, false, null));
       expect(layoutProperties.describeHeight(), 'h=56.0');
       expect(layoutProperties.describeWidth(), 'w=432.6');
     });