diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/videoedit/VideoEditScreen.kt b/app/src/main/java/com/google/android/samples/socialite/ui/videoedit/VideoEditScreen.kt
index ef6db537..47db5280 100644
--- a/app/src/main/java/com/google/android/samples/socialite/ui/videoedit/VideoEditScreen.kt
+++ b/app/src/main/java/com/google/android/samples/socialite/ui/videoedit/VideoEditScreen.kt
@@ -38,11 +38,11 @@ import androidx.compose.foundation.shape.CircleShape
 import androidx.compose.foundation.shape.RoundedCornerShape
 import androidx.compose.foundation.verticalScroll
 import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.automirrored.filled.VolumeMute
 import androidx.compose.material.icons.filled.Close
 import androidx.compose.material.icons.filled.DonutLarge
 import androidx.compose.material.icons.filled.FormatSize
 import androidx.compose.material.icons.filled.Movie
-import androidx.compose.material.icons.filled.VolumeMute
 import androidx.compose.material3.Button
 import androidx.compose.material3.ButtonDefaults
 import androidx.compose.material3.CircularProgressIndicator
@@ -59,7 +59,6 @@ import androidx.compose.material3.TextFieldDefaults
 import androidx.compose.material3.TopAppBar
 import androidx.compose.material3.TopAppBarDefaults
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.saveable.rememberSaveable
@@ -79,7 +78,6 @@ import androidx.hilt.navigation.compose.hiltViewModel
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import androidx.media3.common.util.UnstableApi
 import androidx.navigation.NavController
-import androidx.navigation.compose.rememberNavController
 import com.google.android.samples.socialite.R
 
 private const val TAG = "VideoEditScreen"
@@ -91,37 +89,67 @@ fun VideoEditScreen(
     uri: String,
     onCloseButtonClicked: () -> Unit,
     navController: NavController,
+    viewModel: VideoEditScreenViewModel = hiltViewModel(),
 ) {
-    val context = LocalContext.current
-
-    val viewModel: VideoEditScreenViewModel = hiltViewModel()
     viewModel.setChatId(chatId)
-
-    val isFinishedEditing = viewModel.isFinishedEditing.collectAsStateWithLifecycle()
-    if (isFinishedEditing.value) {
-        navController.popBackStack("chat/$chatId", false)
-    }
-
-    val isProcessing = viewModel.isProcessing.collectAsState()
-
+    val isFinishedEditing by viewModel.isFinishedEditing.collectAsStateWithLifecycle()
+    val isProcessing by viewModel.isProcessing.collectAsStateWithLifecycle()
     var removeAudioEnabled by rememberSaveable { mutableStateOf(false) }
     var overlayText by rememberSaveable { mutableStateOf("") }
     var redOverlayTextEnabled by rememberSaveable { mutableStateOf(false) }
     var largeOverlayTextEnabled by rememberSaveable { mutableStateOf(false) }
+    val context = LocalContext.current
+    if (isFinishedEditing) {
+        navController.popBackStack("chat/$chatId", false)
+    }
+
+    VideoEditScreen(
+        uri = uri,
+        isProcessing = isProcessing,
+        filterChipSelected = removeAudioEnabled,
+        onCloseButtonClicked = onCloseButtonClicked,
+        buttonOverlayText = overlayText,
+        redOverlayTextEnabled = redOverlayTextEnabled,
+        largeOverlayTextEnabled = largeOverlayTextEnabled,
+        onSendButtonClicked = {
+            viewModel.applyVideoTransformation(
+                context = context,
+                videoUri = uri,
+                removeAudio = removeAudioEnabled,
+                textOverlayText = overlayText,
+                textOverlayRedSelected = redOverlayTextEnabled,
+                textOverlayLargeSelected = largeOverlayTextEnabled,
+            )
+        },
+        onFilterChipPressed = { removeAudioEnabled = !removeAudioEnabled },
+        onButtonOverlayTextChanged = { overlayText = it },
+        onRedOverlayTextCheckedStateChanged = { redOverlayTextEnabled = !redOverlayTextEnabled },
+        onLargeOverlayTextCheckedStateChanged = {
+            largeOverlayTextEnabled = !largeOverlayTextEnabled
+        },
+    )
+}
 
+@androidx.annotation.OptIn(UnstableApi::class)
+@Composable
+fun VideoEditScreen(
+    uri: String,
+    isProcessing: Boolean,
+    filterChipSelected: Boolean,
+    buttonOverlayText: String,
+    redOverlayTextEnabled: Boolean,
+    largeOverlayTextEnabled: Boolean,
+    onCloseButtonClicked: () -> Unit,
+    onSendButtonClicked: () -> Unit,
+    onFilterChipPressed: () -> Unit,
+    onButtonOverlayTextChanged: (String) -> Unit,
+    onRedOverlayTextCheckedStateChanged: () -> Unit,
+    onLargeOverlayTextCheckedStateChanged: () -> Unit,
+) {
     Scaffold(
         topBar = {
             VideoEditTopAppBar(
-                onSendButtonClicked = {
-                    viewModel.applyVideoTransformation(
-                        context = context,
-                        videoUri = uri,
-                        removeAudio = removeAudioEnabled,
-                        textOverlayText = overlayText,
-                        textOverlayRedSelected = redOverlayTextEnabled,
-                        textOverlayLargeSelected = largeOverlayTextEnabled,
-                    )
-                },
+                onSendButtonClicked = onSendButtonClicked,
                 onCloseButtonClicked = onCloseButtonClicked,
             )
         },
@@ -136,7 +164,7 @@ fun VideoEditScreen(
             horizontalAlignment = Alignment.CenterHorizontally,
         ) {
             Spacer(modifier = Modifier.height(50.dp))
-            VideoMessagePreview(uri, isProcessing.value)
+            VideoMessagePreview(uri, isProcessing)
             Spacer(modifier = Modifier.height(20.dp))
 
             Column(
@@ -149,28 +177,24 @@ fun VideoEditScreen(
                     .padding(15.dp),
             ) {
                 VideoEditFilterChip(
-                    icon = Icons.Filled.VolumeMute,
-                    selected = removeAudioEnabled,
-                    onClick = { removeAudioEnabled = !removeAudioEnabled },
+                    icon = Icons.AutoMirrored.Filled.VolumeMute,
+                    selected = filterChipSelected,
+                    onClick = onFilterChipPressed,
                     label = stringResource(id = R.string.remove_audio),
                 )
                 Spacer(modifier = Modifier.height(10.dp))
                 TextOverlayOption(
-                    inputtedText = overlayText,
+                    inputtedText = buttonOverlayText,
                     inputtedTextChange = {
                         // Limit character count to 20
                         if (it.length <= 20) {
-                            overlayText = it
+                            onButtonOverlayTextChanged(it)
                         }
                     },
                     redTextCheckedState = redOverlayTextEnabled,
-                    redTextCheckedStateChange = {
-                        redOverlayTextEnabled = !redOverlayTextEnabled
-                    },
+                    redTextCheckedStateChange = onRedOverlayTextCheckedStateChanged,
                     largeTextCheckedState = largeOverlayTextEnabled,
-                    largeTextCheckedStateChange = {
-                        largeOverlayTextEnabled = !largeOverlayTextEnabled
-                    },
+                    largeTextCheckedStateChange = onLargeOverlayTextCheckedStateChanged,
                 )
             }
         }
@@ -308,7 +332,6 @@ fun TextOverlayOption(
     }
 }
 
-@OptIn(ExperimentalMaterial3Api::class)
 @Composable
 private fun VideoEditFilterChip(
     icon: ImageVector,
@@ -343,9 +366,17 @@ private fun VideoEditFilterChip(
 @Preview
 fun VideoEditScreenPreview() {
     VideoEditScreen(
-        chatId = 0L,
         uri = "",
+        isProcessing = true,
+        buttonOverlayText = "",
+        filterChipSelected = true,
+        redOverlayTextEnabled = false,
+        largeOverlayTextEnabled = false,
         onCloseButtonClicked = {},
-        navController = rememberNavController(),
+        onSendButtonClicked = {},
+        onFilterChipPressed = {},
+        onButtonOverlayTextChanged = {},
+        onRedOverlayTextCheckedStateChanged = {},
+        onLargeOverlayTextCheckedStateChanged = {},
     )
 }