Remove temp workaround for feature flag

Using the trunk stable feature flagging now.

Bug: 295516544
Test: m
Test: manual adb
Test: atest SafetySourceBroadcastReceiverTest
Test: atest PrivateSpaceSafetySourceTest
Change-Id: I3d2d8c11dfb224ad2b09572493633366b4c4805c
diff --git a/src/com/android/settings/privatespace/PrivateSpaceDashboardFragment.java b/src/com/android/settings/privatespace/PrivateSpaceDashboardFragment.java
index 9e1d0d5..f72bcd9 100644
--- a/src/com/android/settings/privatespace/PrivateSpaceDashboardFragment.java
+++ b/src/com/android/settings/privatespace/PrivateSpaceDashboardFragment.java
@@ -18,7 +18,7 @@
 
 import android.app.settings.SettingsEnums;
 import android.content.Context;
-import android.util.FeatureFlagUtils;
+import android.os.Flags;
 
 import com.android.settings.R;
 import com.android.settings.dashboard.DashboardFragment;
@@ -56,13 +56,8 @@
             new BaseSearchIndexProvider(R.xml.private_space_settings) {
                 @Override
                 protected boolean isPageSearchEnabled(Context context) {
-                    // Temporary workaround for hiding PS Settings until the trunk stable feature
-                    // flag is available.
-                    // TODO(b/295516544): Remove this workaround when trunk stable feature flag is
-                    // available.
                     return SafetyCenterManagerWrapper.get().isEnabled(context)
-                            && FeatureFlagUtils.isEnabled(context,
-                            FeatureFlagUtils.SETTINGS_PRIVATE_SPACE_SETTINGS);
+                            && Flags.allowPrivateProfile();
                 }
 
                 @Override
diff --git a/src/com/android/settings/privatespace/PrivateSpaceSafetySource.java b/src/com/android/settings/privatespace/PrivateSpaceSafetySource.java
index b07c623..4910a7b 100644
--- a/src/com/android/settings/privatespace/PrivateSpaceSafetySource.java
+++ b/src/com/android/settings/privatespace/PrivateSpaceSafetySource.java
@@ -20,11 +20,11 @@
 import android.app.settings.SettingsEnums;
 import android.content.Context;
 import android.content.Intent;
+import android.os.Flags;
 import android.os.UserManager;
 import android.safetycenter.SafetyEvent;
 import android.safetycenter.SafetySourceData;
 import android.safetycenter.SafetySourceStatus;
-import android.util.FeatureFlagUtils;
 import android.util.Log;
 
 import com.android.settings.R;
@@ -35,7 +35,7 @@
 /** Private Space safety source for the Safety Center */
 public final class PrivateSpaceSafetySource {
     public static final String SAFETY_SOURCE_ID = "AndroidPrivateSpace";
-    private static final String TAG = "PrivateSpaceSafetySource";
+    private static final String TAG = "PrivateSpaceSafetySrc";
 
     private PrivateSpaceSafetySource() {}
 
@@ -54,10 +54,7 @@
             return;
         }
 
-        // Temporary workaround to help prevent the PS Settings showing up in droidfood builds.
-        // TODO(b/295516544): remove this when the trunk stable feature flag for PS is available.
-        if (!FeatureFlagUtils.isEnabled(context,
-                FeatureFlagUtils.SETTINGS_PRIVATE_SPACE_SETTINGS)) {
+        if (!Flags.allowPrivateProfile()) {
             // Setting null safetySourceData so that an old entry gets cleared out and this way
             // provide a response since SC always expects one on rescan.
             SafetyCenterManagerWrapper.get().setSafetySourceData(
diff --git a/tests/unit/Android.bp b/tests/unit/Android.bp
index cc003ea..a38684c 100644
--- a/tests/unit/Android.bp
+++ b/tests/unit/Android.bp
@@ -32,6 +32,7 @@
         "truth-prebuilt",
         "androidx.test.uiautomator_uiautomator",
         "kotlinx_coroutines_test",
+        "flag-junit-base",
         // Don't add SettingsLib libraries here - you can use them directly as they are in the
         // instrumented Settings app.
     ],
diff --git a/tests/unit/src/com/android/settings/privatespace/PrivateSpaceSafetySourceTest.java b/tests/unit/src/com/android/settings/privatespace/PrivateSpaceSafetySourceTest.java
index 2dc00e1..ddf5287 100644
--- a/tests/unit/src/com/android/settings/privatespace/PrivateSpaceSafetySourceTest.java
+++ b/tests/unit/src/com/android/settings/privatespace/PrivateSpaceSafetySourceTest.java
@@ -27,10 +27,11 @@
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
+import android.os.Flags;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.safetycenter.SafetyEvent;
 import android.safetycenter.SafetySourceData;
 import android.safetycenter.SafetySourceStatus;
-import android.util.FeatureFlagUtils;
 
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
@@ -39,6 +40,7 @@
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -51,15 +53,13 @@
             new SafetyEvent.Builder(SAFETY_EVENT_TYPE_DEVICE_REBOOTED).build();
     private Context mContext = ApplicationProvider.getApplicationContext();
     @Mock private SafetyCenterManagerWrapper mSafetyCenterManagerWrapper;
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
     /** Required setup before a test. */
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         SafetyCenterManagerWrapper.sInstance = mSafetyCenterManagerWrapper;
-
-        FeatureFlagUtils
-                .setEnabled(mContext, FeatureFlagUtils.SETTINGS_PRIVATE_SPACE_SETTINGS, true);
     }
 
     /** Required setup after a test. */
@@ -83,6 +83,7 @@
     @Test
     public void onDeviceRebootedEvent_whenSafetyCenterEnabled_setsData() {
         when(mSafetyCenterManagerWrapper.isEnabled(mContext)).thenReturn(true);
+        mSetFlagsRule.enableFlags(Flags.FLAG_ALLOW_PRIVATE_PROFILE);
 
         PrivateSpaceSafetySource.setSafetySourceData(mContext, EVENT_TYPE_DEVICE_REBOOTED);
 
@@ -90,13 +91,11 @@
                 any(), eq(SAFETY_SOURCE_ID), any(), eq(EVENT_TYPE_DEVICE_REBOOTED));
     }
 
-    // TODO(b/295516544): Modify this test for the new trunk stable flag instead when available.
     /** Tests that when the feature is disabled null data is set. */
     @Test
     public void setSafetySourceData_whenFeatureDisabled_setsNullData() {
         when(mSafetyCenterManagerWrapper.isEnabled(mContext)).thenReturn(true);
-        FeatureFlagUtils
-                .setEnabled(mContext, FeatureFlagUtils.SETTINGS_PRIVATE_SPACE_SETTINGS, false);
+        mSetFlagsRule.disableFlags(Flags.FLAG_ALLOW_PRIVATE_PROFILE);
 
         PrivateSpaceSafetySource.setSafetySourceData(mContext, EVENT_TYPE_DEVICE_REBOOTED);
 
@@ -105,15 +104,13 @@
                 any(), eq(SAFETY_SOURCE_ID), captor.capture(), eq(EVENT_TYPE_DEVICE_REBOOTED));
         SafetySourceData safetySourceData = captor.getValue();
         assertThat(safetySourceData).isNull();
-
-        FeatureFlagUtils
-                .setEnabled(mContext, FeatureFlagUtils.SETTINGS_PRIVATE_SPACE_SETTINGS, true);
     }
 
     /** Tests that setSafetySourceData sets the source status enabled. */
     @Test
     public void setSafetySourceData_setsEnabled() {
         when(mSafetyCenterManagerWrapper.isEnabled(mContext)).thenReturn(true);
+        mSetFlagsRule.enableFlags(Flags.FLAG_ALLOW_PRIVATE_PROFILE);
 
         PrivateSpaceSafetySource.setSafetySourceData(mContext, EVENT_TYPE_DEVICE_REBOOTED);
 
@@ -129,6 +126,7 @@
     @Test
     public void setSafetySourceData_setsPsIntent() {
         when(mSafetyCenterManagerWrapper.isEnabled(mContext)).thenReturn(true);
+        mSetFlagsRule.enableFlags(Flags.FLAG_ALLOW_PRIVATE_PROFILE);
 
         PrivateSpaceSafetySource.setSafetySourceData(mContext, EVENT_TYPE_DEVICE_REBOOTED);
 
diff --git a/tests/unit/src/com/android/settings/safetycenter/SafetySourceBroadcastReceiverTest.java b/tests/unit/src/com/android/settings/safetycenter/SafetySourceBroadcastReceiverTest.java
index caae44a..85bd0e2 100644
--- a/tests/unit/src/com/android/settings/safetycenter/SafetySourceBroadcastReceiverTest.java
+++ b/tests/unit/src/com/android/settings/safetycenter/SafetySourceBroadcastReceiverTest.java
@@ -30,9 +30,10 @@
 
 import android.content.Context;
 import android.content.Intent;
+import android.os.Flags;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.safetycenter.SafetyEvent;
 import android.safetycenter.SafetySourceData;
-import android.util.FeatureFlagUtils;
 
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
@@ -43,6 +44,7 @@
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -62,6 +64,7 @@
     @Mock private SafetyCenterManagerWrapper mSafetyCenterManagerWrapper;
 
     @Mock private LockPatternUtils mLockPatternUtils;
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
     @Before
     public void setUp() {
@@ -241,15 +244,10 @@
     }
 
     /** Tests that the PS source sets null data when it's disabled. */
-    // TODO(b/295516544): Modify this test for the new trunk stable flag instead when available.
     @Test
     public void onReceive_onRefresh_withPrivateSpaceFeatureDisabled_setsNullData() {
         when(mSafetyCenterManagerWrapper.isEnabled(mApplicationContext)).thenReturn(true);
-        FeatureFlagUtils
-                .setEnabled(
-                        mApplicationContext,
-                        FeatureFlagUtils.SETTINGS_PRIVATE_SPACE_SETTINGS,
-                        false);
+        mSetFlagsRule.disableFlags(Flags.FLAG_ALLOW_PRIVATE_PROFILE);
 
         Intent intent =
                 new Intent()
@@ -265,12 +263,6 @@
                 .setSafetySourceData(any(), any(), captor.capture(), any());
 
         assertThat(captor.getValue()).isEqualTo(null);
-
-        FeatureFlagUtils
-                .setEnabled(
-                        mApplicationContext,
-                        FeatureFlagUtils.SETTINGS_PRIVATE_SPACE_SETTINGS,
-                        true);
     }
 
     @Test