packages/wyatt_ui_layout/lib/src/presentation/layouts/structural_layouts/top_app_bar_layout.dart

@@ -59,7 +59,7 @@ abstract class TopBarLayout<T extends TopBarComponent>
   final double height;
 
   /// Returns the top bar component for the given [BuildContext].
-  T appBar(BuildContext context, String? barId);
+  T? appBar(BuildContext context, String? barId);
 
   /// The [scaffoldFieldsWrapper] is a final variable that serves as a wrapper
   /// for customizing the scaffold.
@@ -67,18 +67,21 @@ abstract class TopBarLayout<T extends TopBarComponent>
   /// scaffold functionality.
   final ScaffoldFieldsWrapper? scaffoldFieldsWrapper;
 
+  PreferredSize? _buildAppBar(BuildContext context) {
+    final topAppBar =
+        custom?.call(appBar(context, barId)) ?? appBar(context, barId);
+
+    return topAppBar != null
+        ? PreferredSize(
+            preferredSize: Size.fromHeight(height),
+            child: topAppBar,
+          )
+        : null;
+  }
+
   @override
   Widget build(BuildContext context) => Scaffold(
-        appBar: (custom?.call(
+        appBar: _buildAppBar(context),
-                  appBar(context, barId),
-                ) !=
-                null)
-            ? PreferredSize(
-                preferredSize: Size.fromHeight(height),
-                child: custom?.call(appBar(context, barId)) ??
-                    appBar(context, barId),
-              )
-            : null,
         body: body,
         floatingActionButtonLocation:
             scaffoldFieldsWrapper?.floatingActionButtonLocation,
@@ -130,9 +133,23 @@ class TopAppBarLayout extends TopBarLayout<TopAppBarComponent> {
     super.key,
   });
 
+  PreferredSize? _buildAppBar(BuildContext context) {
+    final appBar = custom?.call(
+          context.components.topAppBarComponentOrNull(barId)?.call(),
+        ) ??
+        context.components.topAppBarComponentOrNull(barId)?.call();
+
+    return appBar != null
+        ? PreferredSize(
+            preferredSize: Size.fromHeight(height),
+            child: appBar,
+          )
+        : null;
+  }
+
   @override
-  TopAppBarComponent appBar(BuildContext context, String? barId) =>
+  TopAppBarComponent? appBar(BuildContext context, String? barId) =>
-      context.components.topAppBarComponent(barId).call();
+      context.components.topAppBarComponentOrNull(barId)?.call();
 }
 
 /// A concrete implementation of [TopBarLayout] for a navigation bar.
@@ -153,6 +170,6 @@ class TopNavigationBarLayout extends TopBarLayout<TopNavigationBarComponent> {
   });
 
   @override
-  TopNavigationBarComponent appBar(BuildContext context, String? barId) =>
+  TopNavigationBarComponent? appBar(BuildContext context, String? barId) =>
-      context.components.topNavigationBarComponent(barId).call();
+      context.components.topNavigationBarComponentOrNull(barId)?.call();
 }