Merge "De-flake scheduler tests." into udc-mainline-prod
diff --git a/apk/res/layout/fragment_data_entries.xml b/apk/res/layout/fragment_data_entries.xml
index 55637ac..a375ae1 100644
--- a/apk/res/layout/fragment_data_entries.xml
+++ b/apk/res/layout/fragment_data_entries.xml
@@ -36,7 +36,8 @@
         android:layout_height="wrap_content"
         android:gravity="center"
         android:text="@string/no_data"
-        android:visibility="gone" />
+        android:visibility="gone"
+        android:textAppearance="?attr/textAppearanceSummaryLarge"/>
 
     <TextView
         android:id="@+id/error_view"
@@ -44,11 +45,11 @@
         android:layout_height="wrap_content"
         android:gravity="center"
         android:text="@string/default_error"
-        android:visibility="gone" />
+        android:visibility="gone"
+        android:textAppearance="?attr/textAppearanceSummaryLarge"/>
 
     <androidx.recyclerview.widget.RecyclerView
         android:id="@+id/data_entries_list"
         android:layout_width="match_parent"
         android:layout_height="wrap_content" />
-    
 </LinearLayout>
diff --git a/apk/res/layout/item_exercise_session_entry.xml b/apk/res/layout/item_exercise_session_entry.xml
index 92480e1..89eda9a 100644
--- a/apk/res/layout/item_exercise_session_entry.xml
+++ b/apk/res/layout/item_exercise_session_entry.xml
@@ -20,7 +20,6 @@
     xmlns:app="http://schemas.android.com/apk/res-auto"
     android:minHeight="?android:attr/listPreferredItemHeightSmall"
     android:background="?android:attr/selectableItemBackground"
-    android:clickable="true"
     android:paddingStart="?android:attr/listPreferredItemPaddingStart"
     android:clipToPadding="false">
 
diff --git a/apk/src/com/android/healthconnect/controller/autodelete/AutoDeleteFragment.kt b/apk/src/com/android/healthconnect/controller/autodelete/AutoDeleteFragment.kt
index 364ff41..3011ca2 100644
--- a/apk/src/com/android/healthconnect/controller/autodelete/AutoDeleteFragment.kt
+++ b/apk/src/com/android/healthconnect/controller/autodelete/AutoDeleteFragment.kt
@@ -43,6 +43,7 @@
         private const val AUTO_DELETE_SECTION = "auto_delete_section"
         private const val HEADER = "header"
     }
+
     init {
         this.setPageName(PageName.AUTO_DELETE_PAGE)
     }
@@ -95,7 +96,12 @@
         }
 
         mHeaderSection?.removeAll()
-        mHeaderSection?.addPreference(HeaderPreference(requireContext(), requireActivity()))
+        mHeaderSection?.addPreference(
+            HeaderPreference(requireContext(), requireActivity()).also {
+                // prevents a11y readers from announcing that the header is selectable
+                // this does not prevent the link from being clickable
+                it.isSelectable = false
+            })
 
         childFragmentManager.setFragmentResultListener(SET_TO_NEVER_EVENT, this) { _, _ ->
             viewModel.updateAutoDeleteRange(AutoDeleteRange.AUTO_DELETE_RANGE_NEVER)
diff --git a/apk/src/com/android/healthconnect/controller/dataentries/ExerciseSessionItemViewBinder.kt b/apk/src/com/android/healthconnect/controller/dataentries/ExerciseSessionItemViewBinder.kt
index 25162e2..8108482 100644
--- a/apk/src/com/android/healthconnect/controller/dataentries/ExerciseSessionItemViewBinder.kt
+++ b/apk/src/com/android/healthconnect/controller/dataentries/ExerciseSessionItemViewBinder.kt
@@ -79,9 +79,13 @@
             logger.logInteraction(DataEntriesElement.DATA_ENTRY_DELETE_BUTTON)
             onDeleteEntryClicked?.onDeleteEntry(data.uuid, data.dataType, index)
         }
-        container.setOnClickListener {
-            logger.logInteraction(DataEntriesElement.EXERCISE_SESSION_ENTRY_BUTTON)
-            onItemClickedListener?.onItemClicked(data.uuid, index)
+        if (showSecondAction) {
+            container.setOnClickListener {
+                logger.logInteraction(DataEntriesElement.EXERCISE_SESSION_ENTRY_BUTTON)
+                onItemClickedListener?.onItemClicked(data.uuid, index)
+            }
+        } else {
+            container.isClickable = false
         }
     }
 }
diff --git a/apk/src/com/android/healthconnect/controller/dataentries/SeriesDataItemViewBinder.kt b/apk/src/com/android/healthconnect/controller/dataentries/SeriesDataItemViewBinder.kt
index bf873fd..3a20e28 100644
--- a/apk/src/com/android/healthconnect/controller/dataentries/SeriesDataItemViewBinder.kt
+++ b/apk/src/com/android/healthconnect/controller/dataentries/SeriesDataItemViewBinder.kt
@@ -68,9 +68,15 @@
             logger.logInteraction(DataEntriesElement.DATA_ENTRY_DELETE_BUTTON)
             onDeleteEntryClicked?.onDeleteEntry(data.uuid, data.dataType, index)
         }
-        container.setOnClickListener {
-            logger.logInteraction(DataEntriesElement.DATA_ENTRY_VIEW)
-            onItemClickedListener?.onItemClicked(data.uuid, index)
+        // if the delete button is not showing, we are in the Entry details screen and this item
+        // should not be clickable
+        if (showSecondAction) {
+            container.setOnClickListener {
+                logger.logInteraction(DataEntriesElement.DATA_ENTRY_VIEW)
+                onItemClickedListener?.onItemClicked(data.uuid, index)
+            }
+        } else {
+            container.isClickable = false
         }
     }
 }
diff --git a/apk/src/com/android/healthconnect/controller/dataentries/SleepSessionItemViewBinder.kt b/apk/src/com/android/healthconnect/controller/dataentries/SleepSessionItemViewBinder.kt
index 55443c0..d4c25ef 100644
--- a/apk/src/com/android/healthconnect/controller/dataentries/SleepSessionItemViewBinder.kt
+++ b/apk/src/com/android/healthconnect/controller/dataentries/SleepSessionItemViewBinder.kt
@@ -71,9 +71,13 @@
             logger.logInteraction(DataEntriesElement.DATA_ENTRY_DELETE_BUTTON)
             onDeleteEntryListenerClicked?.onDeleteEntry(data.uuid, data.dataType, index)
         }
-        container.setOnClickListener {
-            logger.logInteraction(DataEntriesElement.SLEEP_SESSION_ENTRY_BUTTON)
-            onItemClickedListener?.onItemClicked(data.uuid, index)
+        if (showSecondAction) {
+            container.setOnClickListener {
+                logger.logInteraction(DataEntriesElement.SLEEP_SESSION_ENTRY_BUTTON)
+                onItemClickedListener?.onItemClicked(data.uuid, index)
+            }
+        } else {
+            container.isClickable = false
         }
     }
 }
diff --git a/apk/src/com/android/healthconnect/controller/permissions/app/ConnectedAppFragment.kt b/apk/src/com/android/healthconnect/controller/permissions/app/ConnectedAppFragment.kt
index ebb989d..a17f9b1 100644
--- a/apk/src/com/android/healthconnect/controller/permissions/app/ConnectedAppFragment.kt
+++ b/apk/src/com/android/healthconnect/controller/permissions/app/ConnectedAppFragment.kt
@@ -71,6 +71,7 @@
 import com.android.healthconnect.controller.utils.showLoadingDialog
 import com.android.settingslib.widget.AppHeaderPreference
 import com.android.settingslib.widget.FooterPreference
+import com.android.settingslib.widget.OnMainSwitchChangeListener
 import dagger.hilt.android.AndroidEntryPoint
 import javax.inject.Inject
 
@@ -205,24 +206,25 @@
         }
     }
 
-    private fun setupAllowAllPreference() {
-        allowAllPreference?.addOnSwitchChangeListener { preference, grantAll ->
-            if (preference.isPressed || allowAllPreference!!.isAccessibilityFocused()) {
-                if (grantAll) {
-                    val permissionsUpdated = appPermissionViewModel.grantAllPermissions(packageName)
-                    if (!permissionsUpdated) {
-                        preference.isChecked = false
-                        Toast.makeText(requireContext(), R.string.default_error, Toast.LENGTH_SHORT)
-                            .show()
-                    }
-                } else {
-                    showRevokeAllPermissions()
-                }
+    private val onSwitchChangeListener = OnMainSwitchChangeListener { switchView, isChecked ->
+        if (isChecked) {
+            val permissionsUpdated = appPermissionViewModel.grantAllPermissions(packageName)
+            if (!permissionsUpdated) {
+                switchView.isChecked = false
+                Toast.makeText(requireContext(), R.string.default_error, Toast.LENGTH_SHORT).show()
             }
+        } else {
+            showRevokeAllPermissions()
         }
+    }
+
+    private fun setupAllowAllPreference() {
+        allowAllPreference?.addOnSwitchChangeListener(onSwitchChangeListener)
         appPermissionViewModel.allAppPermissionsGranted.observe(viewLifecycleOwner) { isAllGranted
             ->
+            allowAllPreference?.removeOnSwitchChangeListener(onSwitchChangeListener)
             allowAllPreference?.isChecked = isAllGranted
+            allowAllPreference?.addOnSwitchChangeListener(onSwitchChangeListener)
         }
     }
 
@@ -263,7 +265,7 @@
                     }
 
                 val preference =
-                    HealthSwitchPreference(requireContext()).also {
+                    HealthSwitchPreference(requireContext()).also { it ->
                         val healthCategory =
                             fromHealthPermissionType(permission.healthPermissionType)
                         it.icon = healthCategory.icon(requireContext())
@@ -272,6 +274,7 @@
                         it.logNameActive = AppAccessElement.PERMISSION_SWITCH_ACTIVE
                         it.logNameInactive = AppAccessElement.PERMISSION_SWITCH_INACTIVE
                         it.setOnPreferenceChangeListener { _, newValue ->
+                            allowAllPreference?.removeOnSwitchChangeListener(onSwitchChangeListener)
                             val checked = newValue as Boolean
                             val permissionUpdated =
                                 appPermissionViewModel.updatePermission(
@@ -283,6 +286,7 @@
                                         Toast.LENGTH_SHORT)
                                     .show()
                             }
+                            allowAllPreference?.addOnSwitchChangeListener(onSwitchChangeListener)
                             permissionUpdated
                         }
                     }
diff --git a/apk/src/com/android/healthconnect/controller/shared/preference/HealthMainSwitchPreference.kt b/apk/src/com/android/healthconnect/controller/shared/preference/HealthMainSwitchPreference.kt
index 1114ce4..0148c4a 100644
--- a/apk/src/com/android/healthconnect/controller/shared/preference/HealthMainSwitchPreference.kt
+++ b/apk/src/com/android/healthconnect/controller/shared/preference/HealthMainSwitchPreference.kt
@@ -17,13 +17,10 @@
 
 import android.content.Context
 import android.util.AttributeSet
-import androidx.preference.PreferenceViewHolder
-import com.android.healthconnect.controller.R
 import com.android.healthconnect.controller.utils.logging.ElementName
 import com.android.healthconnect.controller.utils.logging.ErrorPageElement
 import com.android.healthconnect.controller.utils.logging.HealthConnectLogger
 import com.android.healthconnect.controller.utils.logging.HealthConnectLoggerEntryPoint
-import com.android.settingslib.widget.MainSwitchBar
 import com.android.settingslib.widget.MainSwitchPreference
 import com.android.settingslib.widget.OnMainSwitchChangeListener
 import dagger.hilt.android.EntryPointAccessors
@@ -37,7 +34,6 @@
     var logNameActive: ElementName = ErrorPageElement.UNKNOWN_ELEMENT
     var logNameInactive: ElementName = ErrorPageElement.UNKNOWN_ELEMENT
     private var loggingSwitchListener: OnMainSwitchChangeListener
-    private var mainSwitchBar: MainSwitchBar? = null
 
     init {
         val hiltEntryPoint =
@@ -70,17 +66,4 @@
             logger.logImpression(logNameInactive)
         }
     }
-
-    override fun onBindViewHolder(holder: PreferenceViewHolder?) {
-        super.onBindViewHolder(holder)
-        this.mainSwitchBar = holder?.findViewById(R.id.settingslib_main_switch_bar) as MainSwitchBar
-    }
-
-    // Whether this switch preference is currently focused by a11y readers
-    fun isAccessibilityFocused(): Boolean {
-        mainSwitchBar?.isAccessibilityFocused?.let {
-            return it
-        }
-        return false
-    }
 }
diff --git a/apk/src/com/android/healthconnect/controller/shared/preference/HealthPreferenceFragment.kt b/apk/src/com/android/healthconnect/controller/shared/preference/HealthPreferenceFragment.kt
index 9a118b5..f78b8dc 100644
--- a/apk/src/com/android/healthconnect/controller/shared/preference/HealthPreferenceFragment.kt
+++ b/apk/src/com/android/healthconnect/controller/shared/preference/HealthPreferenceFragment.kt
@@ -19,6 +19,7 @@
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import android.view.accessibility.AccessibilityEvent
 import android.view.animation.Animation
 import android.view.animation.Animation.AnimationListener
 import android.view.animation.AnimationUtils.loadAnimation
@@ -34,6 +35,7 @@
 import com.android.healthconnect.controller.utils.logging.PageName
 import com.android.healthconnect.controller.utils.logging.ToolbarElement
 import com.android.healthconnect.controller.utils.setupSharedMenu
+import com.google.android.material.appbar.AppBarLayout
 import dagger.hilt.android.EntryPointAccessors
 
 /** A base fragment that represents a page in Health Connect. */
@@ -55,6 +57,9 @@
     override fun onCreate(savedInstanceState: Bundle?) {
         setupLogger()
         super.onCreate(savedInstanceState)
+        val appBarLayout = requireActivity().findViewById<AppBarLayout>(R.id.app_bar)
+        appBarLayout?.importantForAccessibility = View.IMPORTANT_FOR_ACCESSIBILITY_YES
+        appBarLayout?.sendAccessibilityEvent(AccessibilityEvent.TYPE_VIEW_FOCUSED)
     }
 
     override fun onResume() {
@@ -139,7 +144,9 @@
                 animation.setAnimationListener(
                     object : AnimationListener {
                         override fun onAnimationStart(animation: Animation) {}
+
                         override fun onAnimationRepeat(animation: Animation) {}
+
                         override fun onAnimationEnd(animation: Animation) {
                             view.visibility = View.INVISIBLE
                         }
diff --git a/service/java/com/android/server/healthconnect/HealthConnectManagerService.java b/service/java/com/android/server/healthconnect/HealthConnectManagerService.java
index 4715f23..e2ee115 100644
--- a/service/java/com/android/server/healthconnect/HealthConnectManagerService.java
+++ b/service/java/com/android/server/healthconnect/HealthConnectManagerService.java
@@ -40,9 +40,7 @@
 import com.android.server.healthconnect.permission.HealthPermissionIntentAppsTracker;
 import com.android.server.healthconnect.permission.PermissionPackageChangesOrchestrator;
 import com.android.server.healthconnect.storage.TransactionManager;
-import com.android.server.healthconnect.storage.datatypehelpers.AppInfoHelper;
-import com.android.server.healthconnect.storage.datatypehelpers.DeviceInfoHelper;
-import com.android.server.healthconnect.storage.datatypehelpers.HealthDataCategoryPriorityHelper;
+import com.android.server.healthconnect.storage.datatypehelpers.DatabaseHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.MigrationEntityHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.PreferenceHelper;
 
@@ -142,10 +140,7 @@
         }
 
         HealthConnectThreadScheduler.shutdownThreadPools();
-        AppInfoHelper.getInstance().clearCache();
-        DeviceInfoHelper.getInstance().clearCache();
-        HealthDataCategoryPriorityHelper.getInstance().clearCache();
-        PreferenceHelper.getInstance().clearCache();
+        DatabaseHelper.clearAllCache();
         mTransactionManager.onUserSwitching();
         RateLimiter.clearCache();
         HealthConnectThreadScheduler.resetThreadPools();
diff --git a/service/java/com/android/server/healthconnect/HealthConnectServiceImpl.java b/service/java/com/android/server/healthconnect/HealthConnectServiceImpl.java
index a31bebb..553739d 100644
--- a/service/java/com/android/server/healthconnect/HealthConnectServiceImpl.java
+++ b/service/java/com/android/server/healthconnect/HealthConnectServiceImpl.java
@@ -136,6 +136,7 @@
 import com.android.server.healthconnect.storage.datatypehelpers.AppInfoHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.ChangeLogsHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.ChangeLogsRequestHelper;
+import com.android.server.healthconnect.storage.datatypehelpers.DatabaseHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.DeviceInfoHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.HealthDataCategoryPriorityHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.MigrationEntityHelper;
@@ -238,6 +239,10 @@
     @Override
     public void grantHealthPermission(
             @NonNull String packageName, @NonNull String permissionName, @NonNull UserHandle user) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(permissionName);
+        Objects.requireNonNull(user);
+
         throwIllegalStateExceptionIfDataSyncInProgress();
         Trace.traceBegin(TRACE_TAG_GRANT_PERMISSION, TAG_GRANT_PERMISSION);
         mPermissionHelper.grantHealthPermission(packageName, permissionName, user);
@@ -250,6 +255,10 @@
             @NonNull String permissionName,
             @Nullable String reason,
             @NonNull UserHandle user) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(permissionName);
+        Objects.requireNonNull(user);
+
         throwIllegalStateExceptionIfDataSyncInProgress();
         mPermissionHelper.revokeHealthPermission(packageName, permissionName, reason, user);
     }
@@ -257,6 +266,9 @@
     @Override
     public void revokeAllHealthPermissions(
             @NonNull String packageName, @Nullable String reason, @NonNull UserHandle user) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(user);
+
         throwIllegalStateExceptionIfDataSyncInProgress();
         mPermissionHelper.revokeAllHealthPermissions(packageName, reason, user);
     }
@@ -264,6 +276,9 @@
     @Override
     public List<String> getGrantedHealthPermissions(
             @NonNull String packageName, @NonNull UserHandle user) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(user);
+
         throwIllegalStateExceptionIfDataSyncInProgress();
         Trace.traceBegin(TRACE_TAG_READ_PERMISSION, TAG_READ_PERMISSION);
         List<String> grantedPermissions =
@@ -275,6 +290,9 @@
     @Override
     public long getHistoricalAccessStartDateInMilliseconds(
             @NonNull String packageName, @NonNull UserHandle userHandle) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(userHandle);
+
         throwIllegalStateExceptionIfDataSyncInProgress();
         Instant date = mPermissionHelper.getHealthDataStartDateAccess(packageName, userHandle);
         if (date == null) {
@@ -299,6 +317,10 @@
             @NonNull AttributionSource attributionSource,
             @NonNull RecordsParcel recordsParcel,
             @NonNull IInsertRecordsResponseCallback callback) {
+        Objects.requireNonNull(attributionSource);
+        Objects.requireNonNull(recordsParcel);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -408,8 +430,12 @@
      */
     public void aggregateRecords(
             @NonNull AttributionSource attributionSource,
-            AggregateDataRequestParcel request,
-            IAggregateRecordsResponseCallback callback) {
+            @NonNull AggregateDataRequestParcel request,
+            @NonNull IAggregateRecordsResponseCallback callback) {
+        Objects.requireNonNull(attributionSource);
+        Objects.requireNonNull(request);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -501,6 +527,10 @@
             @NonNull AttributionSource attributionSource,
             @NonNull ReadRecordsRequestParcel request,
             @NonNull IReadRecordsResponseCallback callback) {
+        Objects.requireNonNull(attributionSource);
+        Objects.requireNonNull(request);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -687,6 +717,10 @@
             @NonNull AttributionSource attributionSource,
             @NonNull RecordsParcel recordsParcel,
             @NonNull IEmptyResponseCallback callback) {
+        Objects.requireNonNull(attributionSource);
+        Objects.requireNonNull(recordsParcel);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -788,6 +822,10 @@
             @NonNull AttributionSource attributionSource,
             @NonNull ChangeLogTokenRequest request,
             @NonNull IGetChangeLogTokenCallback callback) {
+        Objects.requireNonNull(attributionSource);
+        Objects.requireNonNull(request);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
         final HealthConnectServiceLogger.Builder builder =
@@ -851,7 +889,11 @@
     public void getChangeLogs(
             @NonNull AttributionSource attributionSource,
             @NonNull ChangeLogsRequest token,
-            IChangeLogsResponseCallback callback) {
+            @NonNull IChangeLogsResponseCallback callback) {
+        Objects.requireNonNull(attributionSource);
+        Objects.requireNonNull(token);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
         final HealthConnectServiceLogger.Builder builder =
@@ -964,6 +1006,10 @@
             @NonNull AttributionSource attributionSource,
             @NonNull DeleteUsingFiltersRequestParcel request,
             @NonNull IEmptyResponseCallback callback) {
+        Objects.requireNonNull(attributionSource);
+        Objects.requireNonNull(request);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -1057,6 +1103,10 @@
             @NonNull AttributionSource attributionSource,
             @NonNull DeleteUsingFiltersRequestParcel request,
             @NonNull IEmptyResponseCallback callback) {
+        Objects.requireNonNull(attributionSource);
+        Objects.requireNonNull(request);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -1157,6 +1207,9 @@
             @NonNull String packageName,
             @HealthDataCategory.Type int dataCategory,
             @NonNull IGetPriorityResponseCallback callback) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -1206,6 +1259,10 @@
             @NonNull String packageName,
             @NonNull UpdatePriorityRequestParcel updatePriorityRequest,
             @NonNull IEmptyResponseCallback callback) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(updatePriorityRequest);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -1242,7 +1299,10 @@
 
     @Override
     public void setRecordRetentionPeriodInDays(
-            int days, @NonNull UserHandle user, IEmptyResponseCallback callback) {
+            int days, @NonNull UserHandle user, @NonNull IEmptyResponseCallback callback) {
+        Objects.requireNonNull(user);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -1276,6 +1336,8 @@
 
     @Override
     public int getRecordRetentionPeriodInDays(@NonNull UserHandle user) {
+        Objects.requireNonNull(user);
+
         enforceIsForegroundUser(getCallingUserHandle());
         throwExceptionIfDataSyncInProgress();
         try {
@@ -1301,6 +1363,8 @@
      */
     @Override
     public void getContributorApplicationsInfo(@NonNull IApplicationInfoResponseCallback callback) {
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -1337,6 +1401,8 @@
     /** Retrieves {@link RecordTypeInfoResponse} for each RecordType. */
     @Override
     public void queryAllRecordTypesInfo(@NonNull IRecordTypeInfoResponseCallback callback) {
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -1371,7 +1437,11 @@
      * @see HealthConnectManager#queryAccessLogs
      */
     @Override
-    public void queryAccessLogs(@NonNull String packageName, IAccessLogsResponseCallback callback) {
+    public void queryAccessLogs(
+            @NonNull String packageName, @NonNull IAccessLogsResponseCallback callback) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -1413,7 +1483,10 @@
     @Override
     public void getActivityDates(
             @NonNull ActivityDatesRequestParcel activityDatesRequestParcel,
-            IActivityDatesResponseCallback callback) {
+            @NonNull IActivityDatesResponseCallback callback) {
+        Objects.requireNonNull(activityDatesRequestParcel);
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -1459,7 +1532,10 @@
      *     operation.
      */
     @Override
-    public void startMigration(@NonNull String packageName, IMigrationCallback callback) {
+    public void startMigration(@NonNull String packageName, @NonNull IMigrationCallback callback) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(callback);
+
         int uid = Binder.getCallingUid();
         int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -1503,7 +1579,10 @@
      *     operation.
      */
     @Override
-    public void finishMigration(@NonNull String packageName, IMigrationCallback callback) {
+    public void finishMigration(@NonNull String packageName, @NonNull IMigrationCallback callback) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(callback);
+
         int uid = Binder.getCallingUid();
         int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -1541,8 +1620,12 @@
     @Override
     public void writeMigrationData(
             @NonNull String packageName,
-            MigrationEntityParcel parcel,
-            IMigrationCallback callback) {
+            @NonNull MigrationEntityParcel parcel,
+            @NonNull IMigrationCallback callback) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(parcel);
+        Objects.requireNonNull(callback);
+
         int uid = Binder.getCallingUid();
         int pid = Binder.getCallingPid();
         UserHandle callingUserHandle = getCallingUserHandle();
@@ -1583,7 +1666,12 @@
      *     operation.
      */
     public void insertMinDataMigrationSdkExtensionVersion(
-            @NonNull String packageName, int requiredSdkExtension, IMigrationCallback callback) {
+            @NonNull String packageName,
+            int requiredSdkExtension,
+            @NonNull IMigrationCallback callback) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(callback);
+
         int uid = Binder.getCallingUid();
         int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
@@ -1618,6 +1706,10 @@
             @NonNull StageRemoteDataRequest stageRemoteDataRequest,
             @NonNull UserHandle userHandle,
             @NonNull IDataStagingFinishedCallback callback) {
+        Objects.requireNonNull(stageRemoteDataRequest);
+        Objects.requireNonNull(userHandle);
+        Objects.requireNonNull(callback);
+
         Map<String, ParcelFileDescriptor> origPfdsByFileName =
                 stageRemoteDataRequest.getPfdsByFileName();
         Map<String, HealthConnectException> exceptionsByFileName =
@@ -1681,6 +1773,9 @@
     public void getAllDataForBackup(
             @NonNull StageRemoteDataRequest stageRemoteDataRequest,
             @NonNull UserHandle userHandle) {
+        Objects.requireNonNull(stageRemoteDataRequest);
+        Objects.requireNonNull(userHandle);
+
         mContext.enforceCallingPermission(HEALTH_CONNECT_BACKUP_INTER_AGENT_PERMISSION, null);
         mBackupRestore.getAllDataForBackup(stageRemoteDataRequest, userHandle);
     }
@@ -1699,18 +1794,13 @@
      */
     @Override
     public void deleteAllStagedRemoteData(@NonNull UserHandle userHandle) {
+        Objects.requireNonNull(userHandle);
+
         mContext.enforceCallingPermission(
                 DELETE_STAGED_HEALTH_CONNECT_REMOTE_DATA_PERMISSION, null);
         mBackupRestore.deleteAndResetEverything(userHandle);
         mMigrationStateManager.clearCaches(mContext);
-        AppInfoHelper.getInstance().clearData(mTransactionManager);
-        ActivityDateHelper.getInstance().clearData(mTransactionManager);
-        MigrationEntityHelper.getInstance().clearData(mTransactionManager);
-        HealthDataCategoryPriorityHelper.getInstance().clearData(mTransactionManager);
-        PriorityMigrationHelper.getInstance().clearData(mTransactionManager);
-        ChangeLogsHelper.getInstance().clearData(mTransactionManager);
-        ChangeLogsRequestHelper.getInstance().clearData(mTransactionManager);
-        AccessLogsHelper.getInstance().clearData(mTransactionManager);
+        DatabaseHelper.clearAllData(mTransactionManager);
         RateLimiter.clearCache();
         String[] packageNames = mContext.getPackageManager().getPackagesForUid(getCallingUid());
         for (String packageName : packageNames) {
@@ -1734,6 +1824,8 @@
      */
     @Override
     public void getHealthConnectDataState(@NonNull IGetHealthConnectDataStateCallback callback) {
+        Objects.requireNonNull(callback);
+
         try {
             mDataPermissionEnforcer.enforceAnyOfPermissions(
                     MANAGE_HEALTH_DATA_PERMISSION, Manifest.permission.MIGRATE_HEALTH_CONNECT_DATA);
@@ -1796,6 +1888,8 @@
     @Override
     public void getHealthConnectMigrationUiState(
             @NonNull IGetHealthConnectMigrationUiStateCallback callback) {
+        Objects.requireNonNull(callback);
+
         final int uid = Binder.getCallingUid();
         final int pid = Binder.getCallingPid();
         final UserHandle userHandle = Binder.getCallingUserHandle();
diff --git a/service/java/com/android/server/healthconnect/migration/MigrationStateChangeJob.java b/service/java/com/android/server/healthconnect/migration/MigrationStateChangeJob.java
index 3ed80d6..4827f55 100644
--- a/service/java/com/android/server/healthconnect/migration/MigrationStateChangeJob.java
+++ b/service/java/com/android/server/healthconnect/migration/MigrationStateChangeJob.java
@@ -51,12 +51,11 @@
  */
 public final class MigrationStateChangeJob {
     static final int MIN_JOB_ID = MigrationStateChangeJob.class.hashCode();
-    private static final HealthConnectDeviceConfigManager sHealthConnectDeviceConfigManager =
-            HealthConnectDeviceConfigManager.getInitialisedInstance();
 
     public static void scheduleMigrationCompletionJob(Context context, int userId) {
-        if (!HealthConnectDeviceConfigManager.getInitialisedInstance()
-                .isCompleteStateChangeJobEnabled()) {
+        HealthConnectDeviceConfigManager deviceConfigManager =
+                HealthConnectDeviceConfigManager.getInitialisedInstance();
+        if (!deviceConfigManager.isCompleteStateChangeJobEnabled()) {
             return;
         }
         ComponentName componentName = new ComponentName(context, HealthConnectDailyService.class);
@@ -65,9 +64,7 @@
         extras.putString(EXTRA_JOB_NAME_KEY, MIGRATION_COMPLETE_JOB_NAME);
         JobInfo.Builder builder =
                 new JobInfo.Builder(MIN_JOB_ID + userId, componentName)
-                        .setPeriodic(
-                                sHealthConnectDeviceConfigManager
-                                        .getMigrationCompletionJobRunInterval())
+                        .setPeriodic(deviceConfigManager.getMigrationCompletionJobRunInterval())
                         .setExtras(extras);
 
         HealthConnectDailyService.schedule(
@@ -78,8 +75,9 @@
     }
 
     public static void scheduleMigrationPauseJob(Context context, int userId) {
-        if (!HealthConnectDeviceConfigManager.getInitialisedInstance()
-                .isPauseStateChangeJobEnabled()) {
+        HealthConnectDeviceConfigManager deviceConfigManager =
+                HealthConnectDeviceConfigManager.getInitialisedInstance();
+        if (!deviceConfigManager.isPauseStateChangeJobEnabled()) {
             return;
         }
         ComponentName componentName = new ComponentName(context, HealthConnectDailyService.class);
@@ -88,8 +86,7 @@
         extras.putString(EXTRA_JOB_NAME_KEY, MIGRATION_PAUSE_JOB_NAME);
         JobInfo.Builder builder =
                 new JobInfo.Builder(MIN_JOB_ID + userId, componentName)
-                        .setPeriodic(
-                                sHealthConnectDeviceConfigManager.getMigrationPauseJobRunInterval())
+                        .setPeriodic(deviceConfigManager.getMigrationPauseJobRunInterval())
                         .setExtras(extras);
         HealthConnectDailyService.schedule(
                 Objects.requireNonNull(context.getSystemService(JobScheduler.class))
@@ -100,8 +97,9 @@
 
     /** Execute migration completion job */
     public static void executeMigrationCompletionJob(@NonNull Context context) {
-        if (!HealthConnectDeviceConfigManager.getInitialisedInstance()
-                .isCompleteStateChangeJobEnabled()) {
+        HealthConnectDeviceConfigManager deviceConfigManager =
+                HealthConnectDeviceConfigManager.getInitialisedInstance();
+        if (!deviceConfigManager.isCompleteStateChangeJobEnabled()) {
             return;
         }
         if (MigrationStateManager.getInitialisedInstance().getMigrationState()
@@ -123,13 +121,11 @@
                         .plusMillis(
                                 MigrationStateManager.getInitialisedInstance().getMigrationState()
                                                 == MIGRATION_STATE_IDLE
-                                        ? sHealthConnectDeviceConfigManager
-                                                .getIdleStateTimeoutPeriod()
-                                                .toMillis()
-                                        : sHealthConnectDeviceConfigManager
+                                        ? deviceConfigManager.getIdleStateTimeoutPeriod().toMillis()
+                                        : deviceConfigManager
                                                 .getNonIdleStateTimeoutPeriod()
                                                 .toMillis())
-                        .minusMillis(sHealthConnectDeviceConfigManager.getExecutionTimeBuffer());
+                        .minusMillis(deviceConfigManager.getExecutionTimeBuffer());
 
         if (MigrationStateManager.getInitialisedInstance().getMigrationState()
                         == MIGRATION_STATE_ALLOWED
@@ -140,8 +136,7 @@
             if (!Objects.isNull(allowedStateTimeout)) {
                 Instant parsedAllowedStateTimeout =
                         Instant.parse(allowedStateTimeout)
-                                .minusMillis(
-                                        sHealthConnectDeviceConfigManager.getExecutionTimeBuffer());
+                                .minusMillis(deviceConfigManager.getExecutionTimeBuffer());
                 executionTime =
                         executionTime.isAfter(parsedAllowedStateTimeout)
                                 ? parsedAllowedStateTimeout
@@ -158,8 +153,9 @@
 
     /** Execute migration pausing job. */
     public static void executeMigrationPauseJob(@NonNull Context context) {
-        if (!HealthConnectDeviceConfigManager.getInitialisedInstance()
-                .isPauseStateChangeJobEnabled()) {
+        HealthConnectDeviceConfigManager deviceConfigManager =
+                HealthConnectDeviceConfigManager.getInitialisedInstance();
+        if (!deviceConfigManager.isPauseStateChangeJobEnabled()) {
             return;
         }
         if (MigrationStateManager.getInitialisedInstance().getMigrationState()
@@ -178,10 +174,8 @@
         Instant executionTime =
                 Instant.parse(currentStateStartTime)
                         .plusMillis(
-                                sHealthConnectDeviceConfigManager
-                                        .getInProgressStateTimeoutPeriod()
-                                        .toMillis())
-                        .minusMillis(sHealthConnectDeviceConfigManager.getExecutionTimeBuffer());
+                                deviceConfigManager.getInProgressStateTimeoutPeriod().toMillis())
+                        .minusMillis(deviceConfigManager.getExecutionTimeBuffer());
 
         if (Instant.now().isAfter(executionTime)) {
             // If we move to ALLOWED from IN_PROGRESS, then we have reached the IN_PROGRESS_TIMEOUT
diff --git a/service/java/com/android/server/healthconnect/migration/PriorityMigrationHelper.java b/service/java/com/android/server/healthconnect/migration/PriorityMigrationHelper.java
index 0cb1755..c0d2b4b 100644
--- a/service/java/com/android/server/healthconnect/migration/PriorityMigrationHelper.java
+++ b/service/java/com/android/server/healthconnect/migration/PriorityMigrationHelper.java
@@ -24,12 +24,12 @@
 import android.annotation.NonNull;
 import android.content.ContentValues;
 import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
 import android.health.connect.HealthDataCategory;
 import android.util.Pair;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.healthconnect.storage.TransactionManager;
+import com.android.server.healthconnect.storage.datatypehelpers.DatabaseHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.HealthDataCategoryPriorityHelper;
 import com.android.server.healthconnect.storage.request.CreateTableRequest;
 import com.android.server.healthconnect.storage.request.DeleteTableRequest;
@@ -49,7 +49,7 @@
  *
  * @hide
  */
-public final class PriorityMigrationHelper {
+public final class PriorityMigrationHelper extends DatabaseHelper {
 
     @VisibleForTesting
     public static final String PRE_MIGRATION_TABLE_NAME = "pre_migration_category_priority_table";
@@ -130,19 +130,21 @@
         return new CreateTableRequest(PRE_MIGRATION_TABLE_NAME, getColumnInfo());
     }
 
-    /** Upgrades the database to the latest version. */
-    public void onUpgrade(int oldVersion, int newVersion, @NonNull SQLiteDatabase db) {}
+    @Override
+    protected String getMainTableName() {
+        return PRE_MIGRATION_TABLE_NAME;
+    }
 
     /**
      * Populate the pre-migration priority table if table is newly created by copying entries from
      * priority table.
      */
     private void populatePreMigrationTable() {
-            Map<Integer, List<Long>> existingPriority =
-                    HealthDataCategoryPriorityHelper.getInstance()
-                            .getHealthDataCategoryToAppIdPriorityMapImmutable();
+        Map<Integer, List<Long>> existingPriority =
+                HealthDataCategoryPriorityHelper.getInstance()
+                        .getHealthDataCategoryToAppIdPriorityMapImmutable();
 
-            TransactionManager transactionManager = TransactionManager.getInitialisedInstance();
+        TransactionManager transactionManager = TransactionManager.getInitialisedInstance();
         existingPriority.forEach(
                 (category, priority) -> {
                     if (!priority.isEmpty()) {
@@ -175,8 +177,9 @@
     /**
      * This implementation should return the column names with which the table should be created.
      */
+    @Override
     @NonNull
-    private List<Pair<String, String>> getColumnInfo() {
+    protected List<Pair<String, String>> getColumnInfo() {
         ArrayList<Pair<String, String>> columnInfo = new ArrayList<>();
         columnInfo.add(new Pair<>(CATEGORY_COLUMN_NAME, INTEGER_UNIQUE));
         columnInfo.add(new Pair<>(PRIORITY_ORDER_COLUMN_NAME, TEXT_NOT_NULL));
diff --git a/service/java/com/android/server/healthconnect/storage/HealthConnectDatabase.java b/service/java/com/android/server/healthconnect/storage/HealthConnectDatabase.java
index 8e2a163..e094299 100644
--- a/service/java/com/android/server/healthconnect/storage/HealthConnectDatabase.java
+++ b/service/java/com/android/server/healthconnect/storage/HealthConnectDatabase.java
@@ -28,6 +28,7 @@
 import com.android.server.healthconnect.storage.datatypehelpers.AppInfoHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.ChangeLogsHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.ChangeLogsRequestHelper;
+import com.android.server.healthconnect.storage.datatypehelpers.DatabaseHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.DeviceInfoHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.HealthDataCategoryPriorityHelper;
 import com.android.server.healthconnect.storage.datatypehelpers.MigrationEntityHelper;
@@ -80,16 +81,7 @@
         }
 
         mRecordHelpers.forEach(recordHelper -> recordHelper.onUpgrade(db, oldVersion, newVersion));
-        DeviceInfoHelper.getInstance().onUpgrade(oldVersion, newVersion, db);
-        AppInfoHelper.getInstance().onUpgrade(oldVersion, newVersion, db);
-        ChangeLogsHelper.getInstance().onUpgrade(oldVersion, newVersion, db);
-        ChangeLogsRequestHelper.getInstance().onUpgrade(oldVersion, newVersion, db);
-        HealthDataCategoryPriorityHelper.getInstance().onUpgrade(oldVersion, newVersion, db);
-        ActivityDateHelper.getInstance().onUpgrade(oldVersion, newVersion, db);
-        MigrationEntityHelper.getInstance().onUpgrade(oldVersion, newVersion, db);
-        PriorityMigrationHelper.getInstance().onUpgrade(oldVersion, newVersion, db);
-        PreferenceHelper.getInstance().onUpgrade(oldVersion, newVersion, db);
-        AccessLogsHelper.getInstance().onUpgrade(oldVersion, newVersion, db);
+        DatabaseHelper.onUpgrade(db, oldVersion, newVersion);
     }
 
     @Override
@@ -102,6 +94,7 @@
     @Override
     public void onDowngrade(@NonNull SQLiteDatabase db, int oldVersion, int newVersion) {
         Log.i(TAG, "onDowngrade oldVersion = " + oldVersion + " newVersion = " + newVersion);
+        DatabaseHelper.onDowngrade(db, oldVersion, newVersion);
     }
 
     public File getDatabasePath() {
diff --git a/service/java/com/android/server/healthconnect/storage/datatypehelpers/AccessLogsHelper.java b/service/java/com/android/server/healthconnect/storage/datatypehelpers/AccessLogsHelper.java
index 5fbee9a..7ea56a0 100644
--- a/service/java/com/android/server/healthconnect/storage/datatypehelpers/AccessLogsHelper.java
+++ b/service/java/com/android/server/healthconnect/storage/datatypehelpers/AccessLogsHelper.java
@@ -29,7 +29,6 @@
 import android.annotation.NonNull;
 import android.content.ContentValues;
 import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
 import android.health.connect.accesslog.AccessLog;
 import android.health.connect.accesslog.AccessLog.OperationType;
 import android.health.connect.datatypes.RecordTypeIdentifier;
@@ -52,7 +51,7 @@
  *
  * @hide
  */
-public final class AccessLogsHelper {
+public final class AccessLogsHelper extends DatabaseHelper {
     public static final String TABLE_NAME = "access_logs_table";
     private static final String RECORD_TYPE_COLUMN_NAME = "record_type";
     private static final String APP_ID_COLUMN_NAME = "app_id";
@@ -137,8 +136,9 @@
                                 .toEpochMilli());
     }
 
+    @Override
     @NonNull
-    private List<Pair<String, String>> getColumnInfo() {
+    public List<Pair<String, String>> getColumnInfo() {
         List<Pair<String, String>> columnInfo = new ArrayList<>(NUM_COLS);
         columnInfo.add(new Pair<>(PRIMARY_COLUMN_NAME, PRIMARY_AUTOINCREMENT));
         columnInfo.add(new Pair<>(APP_ID_COLUMN_NAME, INTEGER_NOT_NULL));
@@ -149,7 +149,10 @@
         return columnInfo;
     }
 
-    public void onUpgrade(int oldVersion, int newVersion, SQLiteDatabase db) {}
+    @Override
+    protected String getMainTableName() {
+        return TABLE_NAME;
+    }
 
     public static synchronized AccessLogsHelper getInstance() {
         if (sAccessLogsHelper == null) {
@@ -158,9 +161,4 @@
 
         return sAccessLogsHelper;
     }
-
-    /** Deletes all entries from the database. */
-    public synchronized void clearData(TransactionManager transactionManager) {
-        transactionManager.delete(new DeleteTableRequest(TABLE_NAME));
-    }
 }
diff --git a/service/java/com/android/server/healthconnect/storage/datatypehelpers/ActivityDateHelper.java b/service/java/com/android/server/healthconnect/storage/datatypehelpers/ActivityDateHelper.java
index 9820b07..92db3c4 100644
--- a/service/java/com/android/server/healthconnect/storage/datatypehelpers/ActivityDateHelper.java
+++ b/service/java/com/android/server/healthconnect/storage/datatypehelpers/ActivityDateHelper.java
@@ -23,7 +23,6 @@
 import android.annotation.NonNull;
 import android.content.ContentValues;
 import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
 import android.health.connect.datatypes.Record;
 import android.health.connect.internal.datatypes.RecordInternal;
 import android.health.connect.internal.datatypes.utils.RecordMapper;
@@ -52,7 +51,7 @@
  *
  * @hide
  */
-public final class ActivityDateHelper {
+public final class ActivityDateHelper extends DatabaseHelper {
     private static final String TABLE_NAME = "activity_date_table";
     private static final String EPOCH_DAYS_COLUMN_NAME = "epoch_days";
     private static final String RECORD_TYPE_ID_COLUMN_NAME = "record_type_id";
@@ -70,12 +69,9 @@
                 .addUniqueConstraints(List.of(EPOCH_DAYS_COLUMN_NAME, RECORD_TYPE_ID_COLUMN_NAME));
     }
 
-    /** Called on DB update. */
-    public void onUpgrade(int oldVersion, int newVersion, @NonNull SQLiteDatabase db) {}
-
-    /** Deletes all entries from the database and clears the cache. */
-    public synchronized void clearData(TransactionManager transactionManager) {
-        transactionManager.delete(new DeleteTableRequest(TABLE_NAME));
+    @Override
+    protected String getMainTableName() {
+        return TABLE_NAME;
     }
 
     /** Insert a new activity dates for the given records */
@@ -147,8 +143,9 @@
                 });
     }
 
+    @Override
     @NonNull
-    List<Pair<String, String>> getColumnInfo() {
+    protected List<Pair<String, String>> getColumnInfo() {
         return Arrays.asList(
                 new Pair<>(RecordHelper.PRIMARY_COLUMN_NAME, PRIMARY_AUTOINCREMENT),
                 new Pair<>(EPOCH_DAYS_COLUMN_NAME, INTEGER_NOT_NULL),
diff --git a/service/java/com/android/server/healthconnect/storage/datatypehelpers/AppInfoHelper.java b/service/java/com/android/server/healthconnect/storage/datatypehelpers/AppInfoHelper.java
index 56bd1ae..912987a 100644
--- a/service/java/com/android/server/healthconnect/storage/datatypehelpers/AppInfoHelper.java
+++ b/service/java/com/android/server/healthconnect/storage/datatypehelpers/AppInfoHelper.java
@@ -40,7 +40,6 @@
 import android.content.pm.PackageManager.ApplicationInfoFlags;
 import android.content.pm.PackageManager.NameNotFoundException;
 import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
 import android.graphics.Bitmap;
 import android.graphics.BitmapFactory;
 import android.graphics.Canvas;
@@ -56,7 +55,6 @@
 
 import com.android.server.healthconnect.storage.TransactionManager;
 import com.android.server.healthconnect.storage.request.CreateTableRequest;
-import com.android.server.healthconnect.storage.request.DeleteTableRequest;
 import com.android.server.healthconnect.storage.request.ReadTableRequest;
 import com.android.server.healthconnect.storage.request.UpsertTableRequest;
 import com.android.server.healthconnect.storage.utils.WhereClauses;
@@ -82,7 +80,7 @@
  *
  * @hide
  */
-public final class AppInfoHelper {
+public final class AppInfoHelper extends DatabaseHelper {
     public static final String TABLE_NAME = "application_info_table";
     public static final String APPLICATION_COLUMN_NAME = "app_name";
     public static final String PACKAGE_COLUMN_NAME = "package_name";
@@ -110,17 +108,17 @@
 
     private AppInfoHelper() {}
 
-    /** Deletes all entries from the database and clears the cache. */
-    public synchronized void clearData(TransactionManager transactionManager) {
-        transactionManager.delete(new DeleteTableRequest(TABLE_NAME));
-        clearCache();
-    }
-
+    @Override
     public synchronized void clearCache() {
         mAppInfoMap = null;
         mIdPackageNameMap = null;
     }
 
+    @Override
+    protected String getMainTableName() {
+        return TABLE_NAME;
+    }
+
     /**
      * Returns a requests representing the tables that should be created corresponding to this
      * helper
@@ -213,12 +211,6 @@
     }
 
     /**
-     * Called when a db update happens to make any required changes in appInfoHelper respecting
-     * version upgrade.
-     */
-    public void onUpgrade(int oldVersion, int newVersion, @NonNull SQLiteDatabase db) {}
-
-    /**
      * @return id of {@code packageName} or {@link Constants#DEFAULT_LONG} if the id is not found
      */
     public long getAppInfoId(String packageName) {
@@ -227,7 +219,6 @@
         if (appInfo == null) {
             return DEFAULT_LONG;
         }
-
         return appInfo.getId();
     }
 
@@ -681,8 +672,9 @@
      *
      * <p>PLEASE DON'T USE THIS METHOD TO ADD NEW COLUMNS
      */
+    @Override
     @NonNull
-    private List<Pair<String, String>> getColumnInfo() {
+    protected List<Pair<String, String>> getColumnInfo() {
         ArrayList<Pair<String, String>> columnInfo = new ArrayList<>();
         columnInfo.add(new Pair<>(RecordHelper.PRIMARY_COLUMN_NAME, PRIMARY));
         columnInfo.add(new Pair<>(PACKAGE_COLUMN_NAME, TEXT_NOT_NULL_UNIQUE));
diff --git a/service/java/com/android/server/healthconnect/storage/datatypehelpers/ChangeLogsHelper.java b/service/java/com/android/server/healthconnect/storage/datatypehelpers/ChangeLogsHelper.java
index 06e885c..4407bde 100644
--- a/service/java/com/android/server/healthconnect/storage/datatypehelpers/ChangeLogsHelper.java
+++ b/service/java/com/android/server/healthconnect/storage/datatypehelpers/ChangeLogsHelper.java
@@ -31,7 +31,6 @@
 import android.annotation.NonNull;
 import android.content.ContentValues;
 import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
 import android.health.connect.accesslog.AccessLog.OperationType;
 import android.health.connect.changelog.ChangeLogsRequest;
 import android.health.connect.changelog.ChangeLogsResponse.DeletedLog;
@@ -62,7 +61,7 @@
  *
  * @hide
  */
-public final class ChangeLogsHelper {
+public final class ChangeLogsHelper extends DatabaseHelper {
     public static final String TABLE_NAME = "change_logs_table";
     private static final String RECORD_TYPE_COLUMN_NAME = "record_type";
     private static final String APP_ID_COLUMN_NAME = "app_id";
@@ -91,8 +90,10 @@
                 .createIndexOn(APP_ID_COLUMN_NAME);
     }
 
-    // Called on DB update.
-    public void onUpgrade(int oldVersion, int newVersion, @NonNull SQLiteDatabase db) {}
+    @Override
+    protected String getMainTableName() {
+        return TABLE_NAME;
+    }
 
     /** Returns change logs post the time when {@code changeLogTokenRequest} was generated */
     public ChangeLogsResponse getChangeLogs(
@@ -165,7 +166,7 @@
     }
 
     @NonNull
-    private List<Pair<String, String>> getColumnInfo() {
+    protected List<Pair<String, String>> getColumnInfo() {
         List<Pair<String, String>> columnInfo = new ArrayList<>(NUM_COLS);
         columnInfo.add(new Pair<>(PRIMARY_COLUMN_NAME, PRIMARY_AUTOINCREMENT));
         columnInfo.add(new Pair<>(RECORD_TYPE_COLUMN_NAME, INTEGER));
@@ -214,11 +215,6 @@
         return new ArrayMap<>(0);
     }
 
-    /** Deletes all entries from the database. */
-    public synchronized void clearData(TransactionManager transactionManager) {
-        transactionManager.delete(new DeleteTableRequest(TABLE_NAME));
-    }
-
     public static final class ChangeLogs {
         private final Map<RecordTypeAndAppIdPair, List<UUID>> mRecordTypeAndAppIdToUUIDMap =
                 new ArrayMap<>();
diff --git a/service/java/com/android/server/healthconnect/storage/datatypehelpers/ChangeLogsRequestHelper.java b/service/java/com/android/server/healthconnect/storage/datatypehelpers/ChangeLogsRequestHelper.java
index 624e6eb..24f573e 100644
--- a/service/java/com/android/server/healthconnect/storage/datatypehelpers/ChangeLogsRequestHelper.java
+++ b/service/java/com/android/server/healthconnect/storage/datatypehelpers/ChangeLogsRequestHelper.java
@@ -30,7 +30,6 @@
 import android.annotation.NonNull;
 import android.content.ContentValues;
 import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
 import android.health.connect.changelog.ChangeLogTokenRequest;
 import android.util.Pair;
 
@@ -56,7 +55,7 @@
  *
  * @hide
  */
-public final class ChangeLogsRequestHelper {
+public final class ChangeLogsRequestHelper extends DatabaseHelper {
     static final int DEFAULT_CHANGE_LOG_TIME_PERIOD_IN_DAYS = 32;
     private static final String TABLE_NAME = "change_log_request_table";
     private static final String PACKAGES_TO_FILTERS_COLUMN_NAME = "packages_to_filter";
@@ -68,12 +67,14 @@
 
     private ChangeLogsRequestHelper() {}
 
-    // Called on DB update.
-    public void onUpgrade(int oldVersion, int newVersion, @NonNull SQLiteDatabase db) {}
+    @Override
+    protected String getMainTableName() {
+        return TABLE_NAME;
+    }
 
     @NonNull
     public CreateTableRequest getCreateTableRequest() {
-        return new CreateTableRequest(TABLE_NAME, getColumnsInfo());
+        return new CreateTableRequest(TABLE_NAME, getColumnInfo());
     }
 
     @NonNull
@@ -112,8 +113,9 @@
                                 .toEpochMilli());
     }
 
+    @Override
     @NonNull
-    private List<Pair<String, String>> getColumnsInfo() {
+    protected List<Pair<String, String>> getColumnInfo() {
         List<Pair<String, String>> columnInfo = new ArrayList<>();
         columnInfo.add(new Pair<>(PRIMARY_COLUMN_NAME, PRIMARY));
         columnInfo.add(new Pair<>(PACKAGES_TO_FILTERS_COLUMN_NAME, TEXT_NOT_NULL));
@@ -175,11 +177,6 @@
                         .insert(new UpsertTableRequest(TABLE_NAME, contentValues)));
     }
 
-    /** Deletes all entries from the database. */
-    public synchronized void clearData(TransactionManager transactionManager) {
-        transactionManager.delete(new DeleteTableRequest(TABLE_NAME));
-    }
-
     /** A class to represent the request corresponding to a token */
     public static final class TokenRequest {
         private final List<String> mPackageNamesToFilter;
diff --git a/service/java/com/android/server/healthconnect/storage/datatypehelpers/DatabaseHelper.java b/service/java/com/android/server/healthconnect/storage/datatypehelpers/DatabaseHelper.java
new file mode 100644
index 0000000..e8f3aa9
--- /dev/null
+++ b/service/java/com/android/server/healthconnect/storage/datatypehelpers/DatabaseHelper.java
@@ -0,0 +1,83 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.healthconnect.storage.datatypehelpers;
+
+import android.annotation.NonNull;
+import android.database.sqlite.SQLiteDatabase;
+import android.util.Pair;
+
+import com.android.server.healthconnect.storage.TransactionManager;
+import com.android.server.healthconnect.storage.request.DeleteTableRequest;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+/**
+ * Parent class for the database helper classes containing common methods
+ *
+ * @hide
+ */
+public abstract class DatabaseHelper {
+
+    private static Set<DatabaseHelper> sDatabaseHelpers = new HashSet<>();
+
+    protected DatabaseHelper() {
+        sDatabaseHelpers.add(this);
+    }
+
+    public static void onUpgrade(@NonNull SQLiteDatabase db, int oldVersion, int newVersion) {
+        for (DatabaseHelper databaseHelper : sDatabaseHelpers) {
+            databaseHelper.onUpgrade(oldVersion, newVersion, db);
+        }
+    }
+
+    public static void onDowngrade(@NonNull SQLiteDatabase db, int oldVersion, int newVersion) {
+        for (DatabaseHelper databaseHelper : sDatabaseHelpers) {
+            databaseHelper.onDowngrade(oldVersion, newVersion, db);
+        }
+    }
+
+    /** Deletes all entries from the database for the helper class and clears the cache. */
+    public static void clearAllData(@NonNull TransactionManager transactionManager) {
+        for (DatabaseHelper databaseHelper : sDatabaseHelpers) {
+            databaseHelper.clearData(transactionManager);
+        }
+        clearAllCache();
+    }
+
+    public static void clearAllCache() {
+        for (DatabaseHelper databaseHelper : sDatabaseHelpers) {
+            databaseHelper.clearCache();
+        }
+    }
+
+    /** Upgrades the database to the latest version. */
+    protected void onUpgrade(int oldVersion, int newVersion, @NonNull SQLiteDatabase db) {}
+
+    protected void onDowngrade(int oldVersion, int newVersion, @NonNull SQLiteDatabase db) {}
+
+    protected void clearData(@NonNull TransactionManager transactionManager) {
+        transactionManager.delete(new DeleteTableRequest(getMainTableName()));
+    }
+
+    protected void clearCache() {}
+
+    protected abstract String getMainTableName();
+
+    protected abstract List<Pair<String, String>> getColumnInfo();
+}
diff --git a/service/java/com/android/server/healthconnect/storage/datatypehelpers/DeviceInfoHelper.java b/service/java/com/android/server/healthconnect/storage/datatypehelpers/DeviceInfoHelper.java
index fe59138..47ab1bb 100644
--- a/service/java/com/android/server/healthconnect/storage/datatypehelpers/DeviceInfoHelper.java
+++ b/service/java/com/android/server/healthconnect/storage/datatypehelpers/DeviceInfoHelper.java
@@ -28,7 +28,6 @@
 import android.annotation.NonNull;
 import android.content.ContentValues;
 import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
 import android.health.connect.datatypes.Device.DeviceType;
 import android.health.connect.internal.datatypes.RecordInternal;
 import android.util.Pair;
@@ -51,7 +50,7 @@
  *
  * @hide
  */
-public class DeviceInfoHelper {
+public class DeviceInfoHelper extends DatabaseHelper {
     private static final String TABLE_NAME = "device_info_table";
     private static final String MANUFACTURER_COLUMN_NAME = "manufacturer";
     private static final String MODEL_COLUMN_NAME = "model";
@@ -103,16 +102,17 @@
         }
     }
 
-    // Called on DB update.
-    public void onUpgrade(int oldVersion, int newVersion, @NonNull SQLiteDatabase db) {
-        // empty by default
-    }
-
+    @Override
     public synchronized void clearCache() {
         mDeviceInfoMap = null;
         mIdDeviceInfoMap = null;
     }
 
+    @Override
+    protected String getMainTableName() {
+        return TABLE_NAME;
+    }
+
     private synchronized void populateDeviceInfoMap() {
         if (mDeviceInfoMap != null) {
             return;
@@ -192,7 +192,7 @@
      * <p>PLEASE DON'T USE THIS METHOD TO ADD NEW COLUMNS
      */
     @NonNull
-    private List<Pair<String, String>> getColumnInfo() {
+    protected List<Pair<String, String>> getColumnInfo() {
         ArrayList<Pair<String, String>> columnInfo = new ArrayList<>();
         columnInfo.add(new Pair<>(RecordHelper.PRIMARY_COLUMN_NAME, PRIMARY));
         columnInfo.add(new Pair<>(MANUFACTURER_COLUMN_NAME, TEXT_NULL));
diff --git a/service/java/com/android/server/healthconnect/storage/datatypehelpers/HealthDataCategoryPriorityHelper.java b/service/java/com/android/server/healthconnect/storage/datatypehelpers/HealthDataCategoryPriorityHelper.java
index d39bc1d..33f0743 100644
--- a/service/java/com/android/server/healthconnect/storage/datatypehelpers/HealthDataCategoryPriorityHelper.java
+++ b/service/java/com/android/server/healthconnect/storage/datatypehelpers/HealthDataCategoryPriorityHelper.java
@@ -28,7 +28,6 @@
 import android.content.pm.PackageInfo;
 import android.content.res.Resources;
 import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
 import android.health.connect.HealthConnectManager;
 import android.health.connect.HealthDataCategory;
 import android.health.connect.HealthPermissions;
@@ -61,7 +60,7 @@
  *
  * @hide
  */
-public class HealthDataCategoryPriorityHelper {
+public class HealthDataCategoryPriorityHelper extends DatabaseHelper {
     private static final String TABLE_NAME = "health_data_category_priority_table";
     private static final String HEALTH_DATA_CATEGORY_COLUMN_NAME = "health_data_category";
     public static final List<Pair<String, Integer>> UNIQUE_COLUMN_INFO =
@@ -80,11 +79,6 @@
 
     private HealthDataCategoryPriorityHelper() {}
 
-    // Called on DB update.
-    public void onUpgrade(int oldVersion, int newVersion, @NonNull SQLiteDatabase db) {
-        // empty by default
-    }
-
     /**
      * Returns a requests representing the tables that should be created corresponding to this
      * helper
@@ -210,16 +204,22 @@
                 newPriorityOrder);
     }
 
-    /** Deletes all entries from the database and clears the cache. */
-    public synchronized void clearData(@NonNull TransactionManager transactionManager) {
-        transactionManager.delete(new DeleteTableRequest(TABLE_NAME));
+    @Override
+    protected synchronized void clearData(@NonNull TransactionManager transactionManager) {
         clearCache();
+        super.clearData(transactionManager);
     }
 
+    @Override
     public synchronized void clearCache() {
         mHealthDataCategoryToAppIdPriorityMap = null;
     }
 
+    @Override
+    protected String getMainTableName() {
+        return TABLE_NAME;
+    }
+
     private Map<Integer, List<Long>> getHealthDataCategoryToAppIdPriorityMap() {
         if (mHealthDataCategoryToAppIdPriorityMap == null) {
             populateDataCategoryToAppIdPriorityMap();
@@ -299,7 +299,7 @@
      * <p>PLEASE DON'T USE THIS METHOD TO ADD NEW COLUMNS
      */
     @NonNull
-    private List<Pair<String, String>> getColumnInfo() {
+    protected List<Pair<String, String>> getColumnInfo() {
         ArrayList<Pair<String, String>> columnInfo = new ArrayList<>();
         columnInfo.add(new Pair<>(RecordHelper.PRIMARY_COLUMN_NAME, PRIMARY));
         columnInfo.add(new Pair<>(HEALTH_DATA_CATEGORY_COLUMN_NAME, INTEGER_UNIQUE));
@@ -374,8 +374,8 @@
             String currPerm = packageInfo.requestedPermissions[i];
             if (HealthConnectManager.isHealthPermission(context, currPerm)
                     && ((packageInfo.requestedPermissionsFlags[i]
-                    & PackageInfo.REQUESTED_PERMISSION_GRANTED)
-                    != 0)) {
+                                    & PackageInfo.REQUESTED_PERMISSION_GRANTED)
+                            != 0)) {
                 int dataCategory = HealthPermissions.getHealthDataCategory(currPerm);
                 if (dataCategory != -1) {
                     dataCategoryWithPermission.add(dataCategory);
diff --git a/service/java/com/android/server/healthconnect/storage/datatypehelpers/MigrationEntityHelper.java b/service/java/com/android/server/healthconnect/storage/datatypehelpers/MigrationEntityHelper.java
index 90f790b..b0c4961 100644
--- a/service/java/com/android/server/healthconnect/storage/datatypehelpers/MigrationEntityHelper.java
+++ b/service/java/com/android/server/healthconnect/storage/datatypehelpers/MigrationEntityHelper.java
@@ -23,7 +23,6 @@
 
 import android.annotation.NonNull;
 import android.content.ContentValues;
-import android.database.sqlite.SQLiteDatabase;
 import android.util.Pair;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -32,6 +31,7 @@
 import com.android.server.healthconnect.storage.request.DeleteTableRequest;
 import com.android.server.healthconnect.storage.request.UpsertTableRequest;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
@@ -41,7 +41,7 @@
  *
  * @hide
  */
-public final class MigrationEntityHelper {
+public final class MigrationEntityHelper extends DatabaseHelper {
 
     @VisibleForTesting public static final String TABLE_NAME = "migration_entity_table";
     private static final String COLUMN_ENTITY_ID = "entity_id";
@@ -69,8 +69,19 @@
                         new Pair<>(COLUMN_ENTITY_ID, TEXT_NOT_NULL_UNIQUE)));
     }
 
-    /** Upgrades the database to the latest version. */
-    public void onUpgrade(int oldVersion, int newVersion, @NonNull SQLiteDatabase db) {}
+    @Override
+    protected String getMainTableName() {
+        return TABLE_NAME;
+    }
+
+    @Override
+    protected List<Pair<String, String>> getColumnInfo() {
+        ArrayList<Pair<String, String>> columnInfo = new ArrayList<>();
+        columnInfo.add(new Pair<>(PRIMARY_COLUMN_NAME, PRIMARY));
+        columnInfo.add(new Pair<>(COLUMN_ENTITY_ID, TEXT_NOT_NULL_UNIQUE));
+
+        return columnInfo;
+    }
 
     /** Returns a request to insert the provided {@code entityId}. */
     @NonNull
diff --git a/service/java/com/android/server/healthconnect/storage/datatypehelpers/PreferenceHelper.java b/service/java/com/android/server/healthconnect/storage/datatypehelpers/PreferenceHelper.java
index e03f482..95d8bd6 100644
--- a/service/java/com/android/server/healthconnect/storage/datatypehelpers/PreferenceHelper.java
+++ b/service/java/com/android/server/healthconnect/storage/datatypehelpers/PreferenceHelper.java
@@ -24,7 +24,6 @@
 import android.annotation.Nullable;
 import android.content.ContentValues;
 import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
 import android.util.Pair;
 
 import com.android.server.healthconnect.storage.TransactionManager;
@@ -46,7 +45,7 @@
  *
  * @hide
  */
-public final class PreferenceHelper {
+public final class PreferenceHelper extends DatabaseHelper {
     private static final String TABLE_NAME = "preference_table";
     private static final String KEY_COLUMN_NAME = "key";
     public static final List<Pair<String, Integer>> UNIQUE_COLUMN_INFO =
@@ -98,10 +97,16 @@
         return getPreferences().get(key);
     }
 
+    @Override
     public synchronized void clearCache() {
         mPreferences = null;
     }
 
+    @Override
+    protected String getMainTableName() {
+        return TABLE_NAME;
+    }
+
     /** Fetch preferences into memory. */
     public void initializePreferences() {
         populatePreferences();
@@ -138,8 +143,9 @@
         }
     }
 
+    @Override
     @NonNull
-    private List<Pair<String, String>> getColumnInfo() {
+    protected List<Pair<String, String>> getColumnInfo() {
         ArrayList<Pair<String, String>> columnInfo = new ArrayList<>();
         columnInfo.add(new Pair<>(KEY_COLUMN_NAME, TEXT_NOT_NULL_UNIQUE));
         columnInfo.add(new Pair<>(VALUE_COLUMN_NAME, TEXT_NULL));
@@ -147,8 +153,6 @@
         return columnInfo;
     }
 
-    public void onUpgrade(int oldVersion, int newVersion, SQLiteDatabase db) {}
-
     public static synchronized PreferenceHelper getInstance() {
         if (sPreferenceHelper == null) {
             sPreferenceHelper = new PreferenceHelper();
diff --git a/tests/cts/src/android/healthconnect/cts/ActiveCaloriesBurnedRecordTest.java b/tests/cts/src/android/healthconnect/cts/ActiveCaloriesBurnedRecordTest.java
index f9d16b6..0c5f9ca 100644
--- a/tests/cts/src/android/healthconnect/cts/ActiveCaloriesBurnedRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/ActiveCaloriesBurnedRecordTest.java
@@ -46,7 +46,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -65,13 +64,6 @@
 public class ActiveCaloriesBurnedRecordTest {
     private static final String TAG = "ActiveCaloriesBurnedRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/BasalBodyTemperatureRecordTest.java b/tests/cts/src/android/healthconnect/cts/BasalBodyTemperatureRecordTest.java
index ac579d4..7837dc3 100644
--- a/tests/cts/src/android/healthconnect/cts/BasalBodyTemperatureRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/BasalBodyTemperatureRecordTest.java
@@ -42,7 +42,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -59,13 +58,6 @@
 public class BasalBodyTemperatureRecordTest {
     private static final String TAG = "BasalBodyTemperatureRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/BasalMetabolicRateRecordTest.java b/tests/cts/src/android/healthconnect/cts/BasalMetabolicRateRecordTest.java
index 321a7f0..de7a085 100644
--- a/tests/cts/src/android/healthconnect/cts/BasalMetabolicRateRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/BasalMetabolicRateRecordTest.java
@@ -49,7 +49,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -70,13 +69,6 @@
 public class BasalMetabolicRateRecordTest {
     private static final String TAG = "BasalMetabolicRateRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/BloodGlucoseRecordTest.java b/tests/cts/src/android/healthconnect/cts/BloodGlucoseRecordTest.java
index d555590..3f68381 100644
--- a/tests/cts/src/android/healthconnect/cts/BloodGlucoseRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/BloodGlucoseRecordTest.java
@@ -42,7 +42,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -59,13 +58,6 @@
 public class BloodGlucoseRecordTest {
     private static final String TAG = "BloodGlucoseRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/BloodPressureRecordTest.java b/tests/cts/src/android/healthconnect/cts/BloodPressureRecordTest.java
index f2d818b..cd5a2ee 100644
--- a/tests/cts/src/android/healthconnect/cts/BloodPressureRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/BloodPressureRecordTest.java
@@ -51,7 +51,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -70,13 +69,6 @@
 public class BloodPressureRecordTest {
     private static final String TAG = "BloodPressureRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/BodyFatRecordTest.java b/tests/cts/src/android/healthconnect/cts/BodyFatRecordTest.java
index ba4cc84..d615882 100644
--- a/tests/cts/src/android/healthconnect/cts/BodyFatRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/BodyFatRecordTest.java
@@ -42,7 +42,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -59,13 +58,6 @@
 public class BodyFatRecordTest {
     private static final String TAG = "BodyFatRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/BodyTemperatureRecordTest.java b/tests/cts/src/android/healthconnect/cts/BodyTemperatureRecordTest.java
index b8e7ac1..6ef5212 100644
--- a/tests/cts/src/android/healthconnect/cts/BodyTemperatureRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/BodyTemperatureRecordTest.java
@@ -42,7 +42,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -59,13 +58,6 @@
 public class BodyTemperatureRecordTest {
     private static final String TAG = "BodyTemperatureRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/BodyWaterMassRecordTest.java b/tests/cts/src/android/healthconnect/cts/BodyWaterMassRecordTest.java
index 25ea4cc..60d02d8 100644
--- a/tests/cts/src/android/healthconnect/cts/BodyWaterMassRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/BodyWaterMassRecordTest.java
@@ -42,7 +42,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -60,13 +59,6 @@
     private static final String TAG = "BodyWaterMassRecordTest";
     private static final Instant TIME = Instant.ofEpochMilli((long) 1e9);
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/BoneMassRecordTest.java b/tests/cts/src/android/healthconnect/cts/BoneMassRecordTest.java
index 98ef778..6fe76b1 100644
--- a/tests/cts/src/android/healthconnect/cts/BoneMassRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/BoneMassRecordTest.java
@@ -42,7 +42,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -59,13 +58,6 @@
 public class BoneMassRecordTest {
     private static final String TAG = "BoneMassRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/CervicalMucusRecordTest.java b/tests/cts/src/android/healthconnect/cts/CervicalMucusRecordTest.java
index 609a4ed..bdcd933 100644
--- a/tests/cts/src/android/healthconnect/cts/CervicalMucusRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/CervicalMucusRecordTest.java
@@ -41,7 +41,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -58,13 +57,6 @@
 public class CervicalMucusRecordTest {
     private static final String TAG = "CervicalMucusRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/CyclingPedalingCadenceRecordTest.java b/tests/cts/src/android/healthconnect/cts/CyclingPedalingCadenceRecordTest.java
index 5ecbb4b..0b6ba2d 100644
--- a/tests/cts/src/android/healthconnect/cts/CyclingPedalingCadenceRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/CyclingPedalingCadenceRecordTest.java
@@ -48,7 +48,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -68,13 +67,6 @@
 
     private static final String TAG = "CyclingPedalingCadenceRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/DataMigrationTest.java b/tests/cts/src/android/healthconnect/cts/DataMigrationTest.java
index 7d9011b..cac8620 100644
--- a/tests/cts/src/android/healthconnect/cts/DataMigrationTest.java
+++ b/tests/cts/src/android/healthconnect/cts/DataMigrationTest.java
@@ -25,10 +25,12 @@
 import static android.health.connect.HealthPermissions.WRITE_HEIGHT;
 import static android.health.connect.datatypes.units.Length.fromMeters;
 import static android.health.connect.datatypes.units.Power.fromWatts;
+import static android.healthconnect.cts.TestUtils.runShellCommand;
 
 import static com.android.compatibility.common.util.SystemUtil.runWithShellPermissionIdentity;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
 
 import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeFalse;
@@ -104,6 +106,8 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
 @RunWith(AndroidJUnit4.class)
@@ -118,7 +122,12 @@
     private static final String HEALTH_PERMISSION_PREFIX = "android.permission.health.";
     private static final String APP_NAME = "Test App";
     private static final String APP_NAME_NEW = "Test App 2";
-
+    private static final String ENABLE_PAUSE_STATE_CHANGE_JOBS_FLAG =
+            "enable_pause_state_change_jobs";
+    private static final String ENABLE_COMPLETE_STATE_CHANGE_JOBS_FLAG =
+            "enable_complete_state_change_jobs";
+    private static final String ENABLE_PAUSE_STATE_CHANGE_JOB_DEFAULT_FLAG_VALUE = "true";
+    private static final String ENABLE_COMPLETE_STATE_CHANGE_JOB_DEFAULT_FLAG_VALUE = "false";
     // DEFAULT_PAGE_SIZE should hold the same value as Constants#DEFAULT_PAGE_SIZE
     private static final int DEFAULT_PAGE_SIZE = 1000;
 
@@ -137,6 +146,9 @@
                     mEndTime, ZoneOffset.systemDefault().getRules().getOffset(mEndTime));
     private Context mTargetContext;
     private HealthConnectManager mManager;
+    private String mEnableCompletionJobsBackup;
+    private String mEnablePauseJobsBackup;
+    private String mStateChangeJobAdbDumpsys;
 
     private static <T, E extends RuntimeException> T blockingCall(
             Consumer<OutcomeReceiver<T, E>> action) {
@@ -680,62 +692,115 @@
     }
 
     @Test
-    public void testStartMigrationFromIdleState() {
-        runWithShellPermissionIdentity(
-                () -> {
-                    assertThat(TestUtils.getHealthConnectDataMigrationState())
-                            .isEqualTo(MIGRATION_STATE_IDLE);
-                    TestUtils.startMigration();
-                    assertThat(TestUtils.getHealthConnectDataMigrationState())
-                            .isEqualTo(HealthConnectDataState.MIGRATION_STATE_IN_PROGRESS);
-                    TestUtils.finishMigration();
-                    assertThat(TestUtils.getHealthConnectDataMigrationState())
-                            .isEqualTo(HealthConnectDataState.MIGRATION_STATE_COMPLETE);
-                },
-                Manifest.permission.MIGRATE_HEALTH_CONNECT_DATA);
+    public void testStartMigrationFromIdleState() throws IOException {
+        try {
+            enableStateChangeJobs();
+            runWithShellPermissionIdentity(
+                    () -> {
+                        assertStateChangeJobDoesNotExist();
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(MIGRATION_STATE_IDLE);
+                        TestUtils.startMigration();
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(HealthConnectDataState.MIGRATION_STATE_IN_PROGRESS);
+                        assertStateChangeJobExists();
+                        TestUtils.finishMigration();
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(HealthConnectDataState.MIGRATION_STATE_COMPLETE);
+                        assertStateChangeJobDoesNotExist();
+                    },
+                    Manifest.permission.MIGRATE_HEALTH_CONNECT_DATA);
+        } finally {
+            restoreDeviceConfigs();
+        }
     }
 
     @Test
-    public void testInsertMinDataMigrationSdkExtensionVersion_upgradeRequired() {
+    public void testInsertMinDataMigrationSdkExtensionVersion_upgradeRequired() throws IOException {
         int version = SdkExtensions.getExtensionVersion(Build.VERSION_CODES.UPSIDE_DOWN_CAKE) + 1;
-        runWithShellPermissionIdentity(
-                () -> {
-                    assertThat(TestUtils.getHealthConnectDataMigrationState())
-                            .isEqualTo(MIGRATION_STATE_IDLE);
-                    TestUtils.insertMinDataMigrationSdkExtensionVersion(version);
-                    assertThat(TestUtils.getHealthConnectDataMigrationState())
-                            .isEqualTo(MIGRATION_STATE_MODULE_UPGRADE_REQUIRED);
-                    TestUtils.startMigration();
-                    assertThat(TestUtils.getHealthConnectDataMigrationState())
-                            .isEqualTo(HealthConnectDataState.MIGRATION_STATE_IN_PROGRESS);
-                    TestUtils.finishMigration();
-                    assertThat(TestUtils.getHealthConnectDataMigrationState())
-                            .isEqualTo(HealthConnectDataState.MIGRATION_STATE_COMPLETE);
-                },
-                Manifest.permission.MIGRATE_HEALTH_CONNECT_DATA);
+        try {
+            enableStateChangeJobs();
+            runWithShellPermissionIdentity(
+                    () -> {
+                        assertStateChangeJobDoesNotExist();
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(MIGRATION_STATE_IDLE);
+                        TestUtils.insertMinDataMigrationSdkExtensionVersion(version);
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(MIGRATION_STATE_MODULE_UPGRADE_REQUIRED);
+                        assertStateChangeJobExists();
+                        TestUtils.startMigration();
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(HealthConnectDataState.MIGRATION_STATE_IN_PROGRESS);
+                        assertStateChangeJobExists();
+                        TestUtils.finishMigration();
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(HealthConnectDataState.MIGRATION_STATE_COMPLETE);
+                        assertStateChangeJobDoesNotExist();
+                    },
+                    Manifest.permission.MIGRATE_HEALTH_CONNECT_DATA);
+        } finally {
+            restoreDeviceConfigs();
+        }
     }
 
     @Test
-    public void testInsertMinDataMigrationSdkExtensionVersion_noUpgradeRequired() {
+    public void testInsertMinDataMigrationSdkExtensionVersion_noUpgradeRequired()
+            throws IOException {
         int version = SdkExtensions.getExtensionVersion(Build.VERSION_CODES.UPSIDE_DOWN_CAKE);
-        runWithShellPermissionIdentity(
-                () -> {
-                    assertThat(TestUtils.getHealthConnectDataMigrationState())
-                            .isEqualTo(MIGRATION_STATE_IDLE);
-                    TestUtils.insertMinDataMigrationSdkExtensionVersion(version);
-                    assertThat(TestUtils.getHealthConnectDataMigrationState())
-                            .isEqualTo(MIGRATION_STATE_ALLOWED);
-                    TestUtils.insertMinDataMigrationSdkExtensionVersion(version);
-                    assertThat(TestUtils.getHealthConnectDataMigrationState())
-                            .isEqualTo(MIGRATION_STATE_ALLOWED);
-                    TestUtils.startMigration();
-                    assertThat(TestUtils.getHealthConnectDataMigrationState())
-                            .isEqualTo(HealthConnectDataState.MIGRATION_STATE_IN_PROGRESS);
-                    TestUtils.finishMigration();
-                    assertThat(TestUtils.getHealthConnectDataMigrationState())
-                            .isEqualTo(HealthConnectDataState.MIGRATION_STATE_COMPLETE);
-                },
-                Manifest.permission.MIGRATE_HEALTH_CONNECT_DATA);
+        try {
+            enableStateChangeJobs();
+            runWithShellPermissionIdentity(
+                    () -> {
+                        assertStateChangeJobDoesNotExist();
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(MIGRATION_STATE_IDLE);
+                        TestUtils.insertMinDataMigrationSdkExtensionVersion(version);
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(MIGRATION_STATE_ALLOWED);
+                        assertStateChangeJobExists();
+                        TestUtils.startMigration();
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(HealthConnectDataState.MIGRATION_STATE_IN_PROGRESS);
+                        assertStateChangeJobExists();
+                        TestUtils.finishMigration();
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(HealthConnectDataState.MIGRATION_STATE_COMPLETE);
+                        assertStateChangeJobDoesNotExist();
+                    },
+                    Manifest.permission.MIGRATE_HEALTH_CONNECT_DATA);
+        } finally {
+            restoreDeviceConfigs();
+        }
+    }
+
+    @Test
+    public void testInsertMinDataMigrationSdkExtensionVersion_jobsDisabled() throws IOException {
+        int version = SdkExtensions.getExtensionVersion(Build.VERSION_CODES.UPSIDE_DOWN_CAKE);
+        try {
+            disableStateChangeJobs();
+            runWithShellPermissionIdentity(
+                    () -> {
+                        assertStateChangeJobDoesNotExist();
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(MIGRATION_STATE_IDLE);
+                        TestUtils.insertMinDataMigrationSdkExtensionVersion(version);
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(MIGRATION_STATE_ALLOWED);
+                        assertStateChangeJobDoesNotExist();
+                        TestUtils.startMigration();
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(HealthConnectDataState.MIGRATION_STATE_IN_PROGRESS);
+                        assertStateChangeJobDoesNotExist();
+                        TestUtils.finishMigration();
+                        assertThat(TestUtils.getHealthConnectDataMigrationState())
+                                .isEqualTo(HealthConnectDataState.MIGRATION_STATE_COMPLETE);
+                        assertStateChangeJobDoesNotExist();
+                    },
+                    Manifest.permission.MIGRATE_HEALTH_CONNECT_DATA);
+        } finally {
+            restoreDeviceConfigs();
+        }
     }
 
     @Test
@@ -827,6 +892,95 @@
         mExpect.that(getBitmapBytes(appInfo.getIcon())).isNotEqualTo(iconBytes);
     }
 
+    private boolean checkJobExists() throws IOException {
+        String scheduledStateChangeJobPattern =
+                "JOB class com\\.android\\.server\\.healthconnect\\.migration\\"
+                        + ".MigrationStateChangeJob:\\d+/\\d+: [a-f0-9]+ @class com\\"
+                        + ".android\\.server\\.healthconnect\\.migration\\"
+                        + ".MigrationStateChangeJob@android/com\\.android\\.server\\"
+                        + ".healthconnect\\.HealthConnectDailyService";
+
+        String commandOutput = runShellCommand("dumpsys jobscheduler");
+
+        Pattern regexPattern = Pattern.compile(scheduledStateChangeJobPattern);
+        Matcher matcher = regexPattern.matcher(commandOutput);
+        storeStateChangeJobAdbDumpsysOutput(commandOutput);
+        return matcher.find();
+    }
+
+    private void storeStateChangeJobAdbDumpsysOutput(String commandOutput) {
+        String stateChangeJobPattern = ".*MigrationStateChangeJob.*";
+        Pattern regexPattern = Pattern.compile(stateChangeJobPattern);
+        Matcher matcher = regexPattern.matcher(commandOutput);
+        StringBuilder output = new StringBuilder();
+        while (matcher.find()) {
+            output.append(matcher.group()).append("\n");
+        }
+        mStateChangeJobAdbDumpsys = output.toString();
+    }
+
+    private String getHealthFitnessDeviceConfig(String key, String defaultValue)
+            throws IOException {
+        String value = runShellCommand("device_config get health_fitness " + key);
+        return value.isBlank() ? defaultValue : value.strip();
+    }
+
+    private void enableStateChangeJobs() throws IOException {
+        mEnableCompletionJobsBackup =
+                getHealthFitnessDeviceConfig(
+                        ENABLE_COMPLETE_STATE_CHANGE_JOBS_FLAG,
+                        ENABLE_COMPLETE_STATE_CHANGE_JOB_DEFAULT_FLAG_VALUE);
+        mEnablePauseJobsBackup =
+                getHealthFitnessDeviceConfig(
+                        ENABLE_PAUSE_STATE_CHANGE_JOBS_FLAG,
+                        ENABLE_PAUSE_STATE_CHANGE_JOB_DEFAULT_FLAG_VALUE);
+        setHealthFitnessDeviceConfig(ENABLE_COMPLETE_STATE_CHANGE_JOBS_FLAG, "true");
+        setHealthFitnessDeviceConfig(ENABLE_PAUSE_STATE_CHANGE_JOBS_FLAG, "true");
+    }
+
+    private void disableStateChangeJobs() throws IOException {
+        mEnableCompletionJobsBackup =
+                getHealthFitnessDeviceConfig(
+                        ENABLE_COMPLETE_STATE_CHANGE_JOBS_FLAG,
+                        ENABLE_COMPLETE_STATE_CHANGE_JOB_DEFAULT_FLAG_VALUE);
+        mEnablePauseJobsBackup =
+                getHealthFitnessDeviceConfig(
+                        ENABLE_PAUSE_STATE_CHANGE_JOBS_FLAG,
+                        ENABLE_PAUSE_STATE_CHANGE_JOB_DEFAULT_FLAG_VALUE);
+        setHealthFitnessDeviceConfig(ENABLE_COMPLETE_STATE_CHANGE_JOBS_FLAG, "false");
+        setHealthFitnessDeviceConfig(ENABLE_PAUSE_STATE_CHANGE_JOBS_FLAG, "false");
+    }
+
+    private void restoreDeviceConfigs() throws IOException {
+        setHealthFitnessDeviceConfig(
+                ENABLE_COMPLETE_STATE_CHANGE_JOBS_FLAG, mEnableCompletionJobsBackup);
+        setHealthFitnessDeviceConfig(ENABLE_PAUSE_STATE_CHANGE_JOBS_FLAG, mEnablePauseJobsBackup);
+    }
+
+    private void setHealthFitnessDeviceConfig(String key, String value) throws IOException {
+        runShellCommand("device_config put health_fitness " + key + " " + value);
+    }
+
+    private void assertStateChangeJobExists() throws IOException {
+        boolean checkJobExists = checkJobExists();
+        assertWithMessage(
+                        "Expected to find a scheduled MigrationStateChangeJob but not found. ADB"
+                                + " 'dumpsys jobscheduler' output for MigrationStateChangeJob: \n"
+                                + mStateChangeJobAdbDumpsys)
+                .that(checkJobExists)
+                .isTrue();
+    }
+
+    private void assertStateChangeJobDoesNotExist() throws IOException {
+        boolean checkJobExists = checkJobExists();
+        assertWithMessage(
+                        "Did not expect a scheduled MigrationStateChangeJob but found one. ADB"
+                                + " 'dumpsys jobscheduler' output for MigrationStateChangeJob: \n"
+                                + mStateChangeJobAdbDumpsys)
+                .that(checkJobExists)
+                .isFalse();
+    }
+
     private void migrate(MigrationEntity... entities) {
         DataMigrationTest.<Void, MigrationException>blockingCallWithPermissions(
                 callback -> mManager.startMigration(mOutcomeExecutor, callback),
diff --git a/tests/cts/src/android/healthconnect/cts/DistanceRecordTest.java b/tests/cts/src/android/healthconnect/cts/DistanceRecordTest.java
index f7073c0..d40325f 100644
--- a/tests/cts/src/android/healthconnect/cts/DistanceRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/DistanceRecordTest.java
@@ -46,7 +46,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -66,13 +65,6 @@
 public class DistanceRecordTest {
     private static final String TAG = "DistanceRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/ElevationGainedRecordTest.java b/tests/cts/src/android/healthconnect/cts/ElevationGainedRecordTest.java
index df7aa40..f5ac2fc 100644
--- a/tests/cts/src/android/healthconnect/cts/ElevationGainedRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/ElevationGainedRecordTest.java
@@ -46,7 +46,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -65,13 +64,6 @@
 public class ElevationGainedRecordTest {
     private static final String TAG = "ElevationGainedRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/ExerciseRouteDisabledFeatureTest.java b/tests/cts/src/android/healthconnect/cts/ExerciseRouteDisabledFeatureTest.java
index b124432..db617cb 100644
--- a/tests/cts/src/android/healthconnect/cts/ExerciseRouteDisabledFeatureTest.java
+++ b/tests/cts/src/android/healthconnect/cts/ExerciseRouteDisabledFeatureTest.java
@@ -31,19 +31,11 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 
 import java.util.List;
 
 public class ExerciseRouteDisabledFeatureTest {
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     private final UiAutomation mUiAutomation =
             InstrumentationRegistry.getInstrumentation().getUiAutomation();
 
diff --git a/tests/cts/src/android/healthconnect/cts/ExerciseSessionRecordTest.java b/tests/cts/src/android/healthconnect/cts/ExerciseSessionRecordTest.java
index ad7ca24..9d729bc 100644
--- a/tests/cts/src/android/healthconnect/cts/ExerciseSessionRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/ExerciseSessionRecordTest.java
@@ -49,7 +49,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -69,13 +68,6 @@
         super();
     }
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/FloorsClimbedRecordTest.java b/tests/cts/src/android/healthconnect/cts/FloorsClimbedRecordTest.java
index 71bed0d..d1bf705 100644
--- a/tests/cts/src/android/healthconnect/cts/FloorsClimbedRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/FloorsClimbedRecordTest.java
@@ -43,7 +43,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -60,13 +59,6 @@
 public class FloorsClimbedRecordTest {
     private static final String TAG = "FloorsClimbedRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/HealthConnectChangeLogsTests.java b/tests/cts/src/android/healthconnect/cts/HealthConnectChangeLogsTests.java
index c31ce68..fc053ec 100644
--- a/tests/cts/src/android/healthconnect/cts/HealthConnectChangeLogsTests.java
+++ b/tests/cts/src/android/healthconnect/cts/HealthConnectChangeLogsTests.java
@@ -46,12 +46,6 @@
 @AppModeFull(reason = "HealthConnectManager is not accessible to instant apps")
 @RunWith(AndroidJUnit4.class)
 public class HealthConnectChangeLogsTests {
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
 
     @After
     public void tearDown() throws InterruptedException {
diff --git a/tests/cts/src/android/healthconnect/cts/HealthConnectManagerTest.java b/tests/cts/src/android/healthconnect/cts/HealthConnectManagerTest.java
index 92cf908..5f917a7 100644
--- a/tests/cts/src/android/healthconnect/cts/HealthConnectManagerTest.java
+++ b/tests/cts/src/android/healthconnect/cts/HealthConnectManagerTest.java
@@ -131,13 +131,6 @@
                         .build());
     }
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() {
         TestUtils.deleteAllStagedRemoteData();
diff --git a/tests/cts/src/android/healthconnect/cts/HeartRateRecordTest.java b/tests/cts/src/android/healthconnect/cts/HeartRateRecordTest.java
index b92c862..df8935d 100644
--- a/tests/cts/src/android/healthconnect/cts/HeartRateRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/HeartRateRecordTest.java
@@ -51,7 +51,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -72,13 +71,6 @@
 public class HeartRateRecordTest {
     private static final String TAG = "HeartRateRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/HeartRateVariabilityRmssdRecordTest.java b/tests/cts/src/android/healthconnect/cts/HeartRateVariabilityRmssdRecordTest.java
index 9066f98..60361c6 100644
--- a/tests/cts/src/android/healthconnect/cts/HeartRateVariabilityRmssdRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/HeartRateVariabilityRmssdRecordTest.java
@@ -41,7 +41,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -59,13 +58,6 @@
     private static final String TAG = "HeartRateVariabilityRmssdRecordTest";
     private static final Instant TIME = Instant.ofEpochMilli((long) 1e9);
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/HeightRecordTest.java b/tests/cts/src/android/healthconnect/cts/HeightRecordTest.java
index 912f6fb..ddf5a81 100644
--- a/tests/cts/src/android/healthconnect/cts/HeightRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/HeightRecordTest.java
@@ -48,7 +48,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -67,13 +66,6 @@
 public class HeightRecordTest {
     private static final String TAG = "HeightRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/HydrationRecordTest.java b/tests/cts/src/android/healthconnect/cts/HydrationRecordTest.java
index 953acb3..83e6a9e 100644
--- a/tests/cts/src/android/healthconnect/cts/HydrationRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/HydrationRecordTest.java
@@ -45,7 +45,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -63,13 +62,6 @@
 public class HydrationRecordTest {
     private static final String TAG = "HydrationRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/IntermenstrualBleedingRecordTest.java b/tests/cts/src/android/healthconnect/cts/IntermenstrualBleedingRecordTest.java
index 0637a9d..241177f 100644
--- a/tests/cts/src/android/healthconnect/cts/IntermenstrualBleedingRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/IntermenstrualBleedingRecordTest.java
@@ -41,7 +41,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -59,13 +58,6 @@
     private static final String TAG = "IntermenstrualBleedingRecordTest";
     private static final Instant TIME = Instant.ofEpochMilli((long) 1e9);
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/LeanBodyMassRecordTest.java b/tests/cts/src/android/healthconnect/cts/LeanBodyMassRecordTest.java
index ec2e2df..8fad9fb 100644
--- a/tests/cts/src/android/healthconnect/cts/LeanBodyMassRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/LeanBodyMassRecordTest.java
@@ -42,7 +42,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -59,13 +58,6 @@
 public class LeanBodyMassRecordTest {
     private static final String TAG = "LeanBodyMassRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/MenstruationFlowRecordTest.java b/tests/cts/src/android/healthconnect/cts/MenstruationFlowRecordTest.java
index 52bad17..af48e57 100644
--- a/tests/cts/src/android/healthconnect/cts/MenstruationFlowRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/MenstruationFlowRecordTest.java
@@ -41,7 +41,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -58,13 +57,6 @@
 public class MenstruationFlowRecordTest {
     private static final String TAG = "MenstruationFlowRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/MenstruationPeriodRecordTest.java b/tests/cts/src/android/healthconnect/cts/MenstruationPeriodRecordTest.java
index 926e4ec..a0aa236 100644
--- a/tests/cts/src/android/healthconnect/cts/MenstruationPeriodRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/MenstruationPeriodRecordTest.java
@@ -41,7 +41,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -60,13 +59,6 @@
     private static final Instant START_TIME = Instant.ofEpochMilli((long) 1e9);
     private static final Instant END_TIME = Instant.ofEpochMilli((long) 1e10);
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/NutritionRecordTest.java b/tests/cts/src/android/healthconnect/cts/NutritionRecordTest.java
index 0be9253..aca523e 100644
--- a/tests/cts/src/android/healthconnect/cts/NutritionRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/NutritionRecordTest.java
@@ -88,7 +88,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -106,13 +105,6 @@
 public class NutritionRecordTest {
     private static final String TAG = "NutritionRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     private List<AggregationType<Mass>> mMassAggregateTypesList =
             Arrays.asList(
                     BIOTIN_TOTAL,
diff --git a/tests/cts/src/android/healthconnect/cts/OvulationTestRecordTest.java b/tests/cts/src/android/healthconnect/cts/OvulationTestRecordTest.java
index 75fa891..c897b19 100644
--- a/tests/cts/src/android/healthconnect/cts/OvulationTestRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/OvulationTestRecordTest.java
@@ -41,7 +41,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -58,13 +57,6 @@
 public class OvulationTestRecordTest {
     private static final String TAG = "OvulationTestRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/OxygenSaturationRecordTest.java b/tests/cts/src/android/healthconnect/cts/OxygenSaturationRecordTest.java
index 41ffe07..f978789 100644
--- a/tests/cts/src/android/healthconnect/cts/OxygenSaturationRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/OxygenSaturationRecordTest.java
@@ -42,7 +42,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -59,13 +58,6 @@
 public class OxygenSaturationRecordTest {
     private static final String TAG = "OxygenSaturationRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/PowerRecordTest.java b/tests/cts/src/android/healthconnect/cts/PowerRecordTest.java
index 6ab1f87..0797735 100644
--- a/tests/cts/src/android/healthconnect/cts/PowerRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/PowerRecordTest.java
@@ -48,7 +48,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -67,13 +66,6 @@
 public class PowerRecordTest {
     private static final String TAG = "PowerRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/RespiratoryRateRecordTest.java b/tests/cts/src/android/healthconnect/cts/RespiratoryRateRecordTest.java
index c487587..d86c828 100644
--- a/tests/cts/src/android/healthconnect/cts/RespiratoryRateRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/RespiratoryRateRecordTest.java
@@ -41,7 +41,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -58,13 +57,6 @@
 public class RespiratoryRateRecordTest {
     private static final String TAG = "RespiratoryRateRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/RestingHeartRateRecordTest.java b/tests/cts/src/android/healthconnect/cts/RestingHeartRateRecordTest.java
index 06d6f45..b537b52 100644
--- a/tests/cts/src/android/healthconnect/cts/RestingHeartRateRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/RestingHeartRateRecordTest.java
@@ -43,7 +43,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -62,13 +61,6 @@
 public class RestingHeartRateRecordTest {
     private static final String TAG = "RestingHeartRateRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/SessionDatatypeDisabledFeatureTest.java b/tests/cts/src/android/healthconnect/cts/SessionDatatypeDisabledFeatureTest.java
index 76662da..44c69f3 100644
--- a/tests/cts/src/android/healthconnect/cts/SessionDatatypeDisabledFeatureTest.java
+++ b/tests/cts/src/android/healthconnect/cts/SessionDatatypeDisabledFeatureTest.java
@@ -33,20 +33,12 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 
 import java.time.Instant;
 import java.util.List;
 
 public class SessionDatatypeDisabledFeatureTest {
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     private final UiAutomation mUiAutomation =
             InstrumentationRegistry.getInstrumentation().getUiAutomation();
 
diff --git a/tests/cts/src/android/healthconnect/cts/SexualActivityRecordTest.java b/tests/cts/src/android/healthconnect/cts/SexualActivityRecordTest.java
index e55dc3c..d63965e 100644
--- a/tests/cts/src/android/healthconnect/cts/SexualActivityRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/SexualActivityRecordTest.java
@@ -41,7 +41,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -58,13 +57,6 @@
 public class SexualActivityRecordTest {
     private static final String TAG = "SexualActivityRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/SleepSessionRecordTest.java b/tests/cts/src/android/healthconnect/cts/SleepSessionRecordTest.java
index a1357f2..d781438 100644
--- a/tests/cts/src/android/healthconnect/cts/SleepSessionRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/SleepSessionRecordTest.java
@@ -43,7 +43,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 
 import java.time.Instant;
@@ -62,13 +61,6 @@
     private static final CharSequence NOTES = "felt sleepy";
     private static final CharSequence TITLE = "Afternoon nap";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/SpeedRecordTest.java b/tests/cts/src/android/healthconnect/cts/SpeedRecordTest.java
index 149b81e..749d1b1 100644
--- a/tests/cts/src/android/healthconnect/cts/SpeedRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/SpeedRecordTest.java
@@ -49,7 +49,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -69,13 +68,6 @@
 
     private static final String TAG = "SpeedRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/StepsCadenceRecordTest.java b/tests/cts/src/android/healthconnect/cts/StepsCadenceRecordTest.java
index d48bac1..efe59d0 100644
--- a/tests/cts/src/android/healthconnect/cts/StepsCadenceRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/StepsCadenceRecordTest.java
@@ -48,7 +48,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -68,13 +67,6 @@
 
     private static final String TAG = "StepsCadenceRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/StepsRecordTest.java b/tests/cts/src/android/healthconnect/cts/StepsRecordTest.java
index f1623ea..ae94512 100644
--- a/tests/cts/src/android/healthconnect/cts/StepsRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/StepsRecordTest.java
@@ -51,7 +51,6 @@
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Assume;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -73,13 +72,6 @@
 public class StepsRecordTest {
     private static final String TAG = "StepsRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/TestUtils.java b/tests/cts/src/android/healthconnect/cts/TestUtils.java
index db5d7aa..a5dddc1 100644
--- a/tests/cts/src/android/healthconnect/cts/TestUtils.java
+++ b/tests/cts/src/android/healthconnect/cts/TestUtils.java
@@ -114,6 +114,7 @@
 import android.health.connect.datatypes.units.Power;
 import android.health.connect.migration.MigrationException;
 import android.os.OutcomeReceiver;
+import android.os.ParcelFileDescriptor;
 import android.util.Log;
 import android.util.Pair;
 
@@ -121,6 +122,11 @@
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.platform.app.InstrumentationRegistry;
 
+import java.io.BufferedReader;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.io.InputStreamReader;
 import java.time.Duration;
 import java.time.Instant;
 import java.time.LocalDate;
@@ -1376,6 +1382,29 @@
                         BODY_MEASUREMENTS, BASAL_METABOLIC_RATE, new ArrayList<>()));
     }
 
+    static String runShellCommand(String command) throws IOException {
+        UiAutomation uiAutomation =
+                androidx.test.platform.app.InstrumentationRegistry.getInstrumentation()
+                        .getUiAutomation();
+        uiAutomation.adoptShellPermissionIdentity();
+        final ParcelFileDescriptor stdout = uiAutomation.executeShellCommand(command);
+        StringBuilder output = new StringBuilder();
+
+        try (BufferedReader reader =
+                new BufferedReader(
+                        new InputStreamReader(new FileInputStream(stdout.getFileDescriptor())))) {
+            char[] buffer = new char[4096];
+            int bytesRead;
+            while ((bytesRead = reader.read(buffer)) != -1) {
+                output.append(buffer, 0, bytesRead);
+            }
+        } catch (FileNotFoundException e) {
+            Log.e(TAG, e.getMessage());
+        }
+
+        return output.toString();
+    }
+
     static final class RecordAndIdentifier {
         private final int id;
         private final Record recordClass;
diff --git a/tests/cts/src/android/healthconnect/cts/TotalCaloriesBurnedRecordTest.java b/tests/cts/src/android/healthconnect/cts/TotalCaloriesBurnedRecordTest.java
index 33a2922..fa4c605 100644
--- a/tests/cts/src/android/healthconnect/cts/TotalCaloriesBurnedRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/TotalCaloriesBurnedRecordTest.java
@@ -44,7 +44,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -64,13 +63,6 @@
 public class TotalCaloriesBurnedRecordTest {
     private static final String TAG = "TotalCaloriesBurnedRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/Vo2MaxRecordTest.java b/tests/cts/src/android/healthconnect/cts/Vo2MaxRecordTest.java
index 983e5d2..e83f2ea 100644
--- a/tests/cts/src/android/healthconnect/cts/Vo2MaxRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/Vo2MaxRecordTest.java
@@ -41,7 +41,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -58,13 +57,6 @@
 public class Vo2MaxRecordTest {
     private static final String TAG = "Vo2MaxRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/WeightRecordTest.java b/tests/cts/src/android/healthconnect/cts/WeightRecordTest.java
index 3cc06ba..c3f7f0a 100644
--- a/tests/cts/src/android/healthconnect/cts/WeightRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/WeightRecordTest.java
@@ -51,7 +51,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -73,13 +72,6 @@
 public class WeightRecordTest {
     private static final String TAG = "WeightRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/cts/src/android/healthconnect/cts/WheelchairPushesRecordTest.java b/tests/cts/src/android/healthconnect/cts/WheelchairPushesRecordTest.java
index a079d64..ed49508 100644
--- a/tests/cts/src/android/healthconnect/cts/WheelchairPushesRecordTest.java
+++ b/tests/cts/src/android/healthconnect/cts/WheelchairPushesRecordTest.java
@@ -43,7 +43,6 @@
 
 import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -60,13 +59,6 @@
 public class WheelchairPushesRecordTest {
     private static final String TAG = "WheelchairPushesRecordTest";
 
-    @Before
-    public void setUp() {
-        // TODO(b/283737434): Update the HC code to use user aware context on permission change.
-        // Temporary fix to set firstGrantTime for the correct user in HSUM.
-        TestUtils.deleteAllStagedRemoteData();
-    }
-
     @After
     public void tearDown() throws InterruptedException {
         TestUtils.verifyDeleteRecords(
diff --git a/tests/integrationtests/src/android/healthconnect/tests/permissions/HealthConnectWithManagePermissionsTest.java b/tests/integrationtests/src/android/healthconnect/tests/permissions/HealthConnectWithManagePermissionsTest.java
index 4cb6c61..eee649d 100644
--- a/tests/integrationtests/src/android/healthconnect/tests/permissions/HealthConnectWithManagePermissionsTest.java
+++ b/tests/integrationtests/src/android/healthconnect/tests/permissions/HealthConnectWithManagePermissionsTest.java
@@ -281,7 +281,7 @@
         // Grant permission
         assertPermNotGrantedForApp(DEFAULT_APP_PACKAGE, DEFAULT_PERM);
         try {
-            grantHealthPermission(DEFAULT_APP_PACKAGE, /* permissionName= */ null);
+            grantHealthPermission(DEFAULT_APP_PACKAGE, DEFAULT_PERM);
             fail("Expected IllegalStateException for data sync in progress.");
         } catch (IllegalStateException exception) {
             assertNotNull(exception);
diff --git a/tests/unittests/src/com/android/server/healthconnect/storage/datatypehelpers/HealthDataCategoryPriorityHelperTest.java b/tests/unittests/src/com/android/server/healthconnect/storage/datatypehelpers/HealthDataCategoryPriorityHelperTest.java
index 24c51a6..54cd620 100644
--- a/tests/unittests/src/com/android/server/healthconnect/storage/datatypehelpers/HealthDataCategoryPriorityHelperTest.java
+++ b/tests/unittests/src/com/android/server/healthconnect/storage/datatypehelpers/HealthDataCategoryPriorityHelperTest.java
@@ -94,7 +94,7 @@
     @After
     public void tearDown() throws Exception {
         mStaticMockSession.finishMocking();
-        mHealthDataCategoryPriorityHelper.clearCache();
+        mHealthDataCategoryPriorityHelper.clearData(mTransactionManager);
     }
 
     @Test