1 files changed, 16 insertions, 31 deletions

diff --git a/src/screens/badge/BadgeSelection.tsx b/src/screens/badge/BadgeSelection.tsx
index ba2e6af4..a67206c4 100644
--- a/src/screens/badge/BadgeSelection.tsx
+++ b/src/screens/badge/BadgeSelection.tsx
@@ -1,5 +1,4 @@
 import {RouteProp} from '@react-navigation/core';
-import {StackNavigationProp} from '@react-navigation/stack';
 import React, {useEffect, useState} from 'react';
 import {Alert, SafeAreaView, StatusBar, StyleSheet, View} from 'react-native';
 import {Text} from 'react-native-animatable';
@@ -12,33 +11,32 @@ import {BACKGROUND_GRADIENT_MAP} from '../../constants';
 import {BADGE_DATA} from '../../constants/badges';
 import {ERROR_BADGES_EXCEED_LIMIT} from '../../constants/strings';
 import {suggestedPeopleBadgesFinished} from '../../store/actions';
-import {BackgroundGradientType, UniversityBadge} from '../../types';
+import {BackgroundGradientType} from '../../types';
 import {SCREEN_HEIGHT, StatusBarHeight} from '../../utils';
 import BadgeList from './BadgeList';
 import BadgeScreenHeader from './BadgeScreenHeader';
+import {useNavigation} from '@react-navigation/native';
 
 /**
  * Home Screen for displaying Tagg Badge Selections
  **/
 
-type BadgeSelectionParamList = {
-  BadgeList: any[];
-};
-
-type BadgeSelectionScreenNavigationProp = StackNavigationProp<
-  BadgeSelectionParamList,
-  'BadgeList'
->;
-
 type BadgeSelectionRouteProp = RouteProp<MainStackParams, 'BadgeSelection'>;
 
 type BadgeSelectionProps = {
-  navigation: BadgeSelectionScreenNavigationProp;
-  route?: BadgeSelectionRouteProp;
+  route: BadgeSelectionRouteProp;
 };
 
-const BadgeSelection: React.FC<BadgeSelectionProps> = ({navigation, route}) => {
+const BadgeSelection: React.FC<BadgeSelectionProps> = ({route}) => {
+  const prevSelectedBadges = route.params.prevSelectedBadges;
+  const [selectedBadges, setSelectedBadges] = useState<string[]>([]);
   const dispatch = useDispatch();
+  const navigation = useNavigation();
+
+  useEffect(() => {
+    setSelectedBadges(prevSelectedBadges);
+  }, [prevSelectedBadges]);
+
   navigation.setOptions({
     headerRight: () => (
       <TouchableOpacity
@@ -46,7 +44,10 @@ const BadgeSelection: React.FC<BadgeSelectionProps> = ({navigation, route}) => {
         onPress={async () => {
           if (route?.params.editing) {
             updateBadgesService(selectedBadges);
-            navigation.goBack();
+            navigation.navigate('UpdateSPPicture', {
+              editing: true,
+              updatedSelectedBadges: selectedBadges,
+            });
           } else {
             if (selectedBadges.length !== 0) {
               const success = await addBadgesService(selectedBadges);
@@ -67,22 +68,6 @@ const BadgeSelection: React.FC<BadgeSelectionProps> = ({navigation, route}) => {
     ),
   });
 
-  const [selectedBadges, setSelectedBadges] = useState<string[]>([]);
-
-  // Get list of badges from the backend and display here
-  useEffect(() => {
-    const extractBadgeNames = (badges: UniversityBadge[]) => {
-      let extractedBadgeNames: string[] = [];
-      badges.forEach((badge) => {
-        extractedBadgeNames.push(badge.name);
-      });
-      setSelectedBadges(extractedBadgeNames);
-    };
-    if (route && route.params.selectedBadges) {
-      extractBadgeNames(route.params.selectedBadges);
-    }
-  }, []);
-
   const selectKey = (key: string) => {
     if (selectedBadges.includes(key)) {
       const selectedBadgesArray = [...selectedBadges];