[4/n] Create first state and handle comp lifecycle

We define the lifecycle for each Compat UI component
using the predicates defind in CompatUISpec.

Define the concept of shared and component specific states.

Flag: com.android.window.flags.app_compat_ui_framework
Bug: 270361630
Fix: 322817374
Test: atest WMShellUnitTests:DefaultCompatUIHandlerTest

Change-Id: I948cc690fd03bd0f16ea0740da9c3a0a5fe24f9a
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIComponent.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIComponent.kt
new file mode 100644
index 0000000..9ee50ac
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIComponent.kt
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.compatui.api
+
+import android.util.Log
+
+/**
+ * The component created after a {@link CompatUISpec} definition
+ */
+class CompatUIComponent(
+    private val spec: CompatUISpec,
+    private val id: String
+) {
+
+    /**
+     * Invoked every time a new CompatUIInfo comes from core
+     * @param newInfo The new CompatUIInfo object
+     * @param sharedState The state shared between all the component
+     */
+    fun update(newInfo: CompatUIInfo, state: CompatUIState) {
+        // TODO(b/322817374): To be removed when the implementation is provided.
+        Log.d("CompatUIComponent", "update() newInfo: $newInfo state:$state")
+    }
+
+    fun release() {
+        // TODO(b/322817374): To be removed when the implementation is provided.
+        Log.d("CompatUIComponent", "release()")
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIComponentIdGenerator.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIComponentIdGenerator.kt
new file mode 100644
index 0000000..7d663fa
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIComponentIdGenerator.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.compatui.api
+
+/**
+ * Any object responsible to generate an id for a component.
+ */
+interface CompatUIComponentIdGenerator {
+
+    /**
+     * Generates the unique id for a component given a {@link CompatUIInfo} and component
+     * {@link CompatUISpec}.
+     * @param compatUIInfo  The object encapsulating information about the current Task.
+     * @param spec  The {@link CompatUISpec} for the component.
+     */
+    fun generateId(compatUIInfo: CompatUIInfo, spec: CompatUISpec): String
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIComponentState.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIComponentState.kt
new file mode 100644
index 0000000..dcaea00
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIComponentState.kt
@@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.compatui.api
+
+/**
+ * Abstraction of all the component specific state. Each
+ * component can create its own state implementing this
+ * tagging interface.
+ */
+interface CompatUIComponentState
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUISharedState.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUISharedState.kt
new file mode 100644
index 0000000..33e0d46
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUISharedState.kt
@@ -0,0 +1,22 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.compatui.api
+
+/**
+ * Represents the state shared between all the components.
+ */
+class CompatUISharedState
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUISpec.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUISpec.kt
index 24c2c8c..a520d5e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUISpec.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUISpec.kt
@@ -17,10 +17,31 @@
 package com.android.wm.shell.compatui.api
 
 /**
+ * Defines the predicates to invoke for understanding if a component can be created or destroyed.
+ */
+class CompatUILifecyclePredicates(
+    // Predicate evaluating to true if the component needs to be created
+    val creationPredicate: (CompatUIInfo, CompatUISharedState) -> Boolean,
+    // Predicate evaluating to true if the component needs to be destroyed
+    val removalPredicate: (
+        CompatUIInfo,
+        CompatUISharedState,
+        CompatUIComponentState?
+    ) -> Boolean,
+    // Builder for the initial state of the component
+    val stateBuilder: (
+        CompatUIInfo,
+        CompatUISharedState
+    ) -> CompatUIComponentState? = { _, _ -> null }
+)
+
+/**
  * Describes each compat ui component to the framework.
  */
-data class CompatUISpec(
+class CompatUISpec(
     // Unique name for the component. It's used for debug and for generating the
     // unique component identifier in the system.
-    val name: String
-)
\ No newline at end of file
+    val name: String,
+    // The lifecycle definition
+    val lifecycle: CompatUILifecyclePredicates
+)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIState.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIState.kt
new file mode 100644
index 0000000..68307b4
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIState.kt
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.compatui.api
+
+/**
+ * Singleton which contains the global state of the compat ui system.
+ */
+class CompatUIState {
+
+    private val components = mutableMapOf<String, CompatUIComponent>()
+
+    val sharedState = CompatUISharedState()
+
+    val componentStates = mutableMapOf<String, CompatUIComponentState>()
+
+    /**
+     * @return The CompatUIComponent for the given componentId if it exists.
+     */
+    fun getUIComponent(componentId: String): CompatUIComponent? =
+        components[componentId]
+
+    /**
+     * Registers a component for a given componentId along with its optional state.
+     * <p/>
+     * @param componentId       The identifier for the component to register.
+     * @param comp              The {@link CompatUIComponent} instance to register.
+     * @param componentState    The optional state specific of the component. Not all components
+     *                          have a specific state so it can be null.
+     */
+    fun registerUIComponent(
+        componentId: String,
+        comp: CompatUIComponent,
+        componentState: CompatUIComponentState?
+    ) {
+        components[componentId] = comp
+        componentState?.let {
+            componentStates[componentId] = componentState
+        }
+    }
+
+    /**
+     * Unregister a component for a given componentId.
+     * <p/>
+     * @param componentId       The identifier for the component to register.
+     */
+    fun unregisterUIComponent(componentId: String) {
+        components.remove(componentId)
+        componentStates.remove(componentId)
+    }
+
+    /**
+     * Get access to the specific {@link CompatUIComponentState} for a {@link CompatUIComponent}
+     * with a given identifier.
+     * <p/>
+     * @param componentId  The identifier of the {@link CompatUIComponent}.
+     * @return The optional state for the component of the provided id.
+     */
+    @Suppress("UNCHECKED_CAST")
+    fun <T : CompatUIComponentState> stateForComponent(componentId: String) =
+        componentStates[componentId] as? T
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/CompatUIEvents.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/CompatUIEvents.kt
index 23205c3..db3fda0 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/CompatUIEvents.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/CompatUIEvents.kt
@@ -27,9 +27,9 @@
 sealed class CompatUIEvents(override val eventId: Int) : CompatUIEvent {
     /** Sent when the size compat restart button appears. */
     data class SizeCompatRestartButtonAppeared(val taskId: Int) :
-            CompatUIEvents(SIZE_COMPAT_RESTART_BUTTON_APPEARED)
+        CompatUIEvents(SIZE_COMPAT_RESTART_BUTTON_APPEARED)
 
     /** Sent when the size compat restart button is clicked. */
     data class SizeCompatRestartButtonClicked(val taskId: Int) :
-            CompatUIEvents(SIZE_COMPAT_RESTART_BUTTON_CLICKED)
-}
\ No newline at end of file
+        CompatUIEvents(SIZE_COMPAT_RESTART_BUTTON_CLICKED)
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/DefaultCompatUIHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/DefaultCompatUIHandler.kt
index 8408ea6..a7d1b42 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/DefaultCompatUIHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/DefaultCompatUIHandler.kt
@@ -16,25 +16,70 @@
 
 package com.android.wm.shell.compatui.impl
 
+import com.android.wm.shell.compatui.api.CompatUIComponent
+import com.android.wm.shell.compatui.api.CompatUIComponentIdGenerator
 import com.android.wm.shell.compatui.api.CompatUIEvent
 import com.android.wm.shell.compatui.api.CompatUIHandler
 import com.android.wm.shell.compatui.api.CompatUIInfo
 import com.android.wm.shell.compatui.api.CompatUIRepository
+import com.android.wm.shell.compatui.api.CompatUIState
 import java.util.function.Consumer
+import java.util.function.IntSupplier
 
 /**
  * Default implementation of {@link CompatUIHandler} to handle CompatUI components
  */
 class DefaultCompatUIHandler(
-    private val compatUIRepository: CompatUIRepository
+    private val compatUIRepository: CompatUIRepository,
+    private val compatUIState: CompatUIState,
+    private val componentIdGenerator: CompatUIComponentIdGenerator
 ) : CompatUIHandler {
 
     private var compatUIEventSender: Consumer<CompatUIEvent>? = null
+
     override fun onCompatInfoChanged(compatUIInfo: CompatUIInfo) {
+        compatUIRepository.iterateOn { spec ->
+            // We get the identifier for the component depending on the task and spec
+            val componentId = componentIdGenerator.generateId(compatUIInfo, spec)
+            // We check in the state if the component already exists
+            var comp = compatUIState.getUIComponent(componentId)
+            if (comp == null) {
+                // We evaluate the predicate
+                if (spec.lifecycle.creationPredicate(compatUIInfo, compatUIState.sharedState)) {
+                    // We create the component and store in the
+                    // global state
+                    comp = CompatUIComponent(spec, componentId)
+                    // We initialize the state for the component
+                    val compState = spec.lifecycle.stateBuilder(
+                        compatUIInfo,
+                        compatUIState.sharedState
+                    )
+                    compatUIState.registerUIComponent(componentId, comp, compState)
+                    // Now we can invoke the update passing the shared state and
+                    // the state specific to the component
+                    comp.update(compatUIInfo, compatUIState)
+                }
+            } else {
+                // The component is present. We check if we need to remove it
+                if (spec.lifecycle.removalPredicate(
+                        compatUIInfo,
+                        compatUIState.sharedState,
+                        compatUIState.stateForComponent(componentId)
+                    )) {
+                    // We clean the component
+                    comp.release()
+                    // We remove the component
+                    compatUIState.unregisterUIComponent(componentId)
+                } else {
+                    // The component exists so we need to invoke the update methods
+                    comp.update(compatUIInfo, compatUIState)
+                }
+            }
+        }
         // Empty at the moment
     }
 
     override fun setCallback(compatUIEventSender: Consumer<CompatUIEvent>?) {
         this.compatUIEventSender = compatUIEventSender
     }
-}
\ No newline at end of file
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/DefaultComponentIdGenerator.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/DefaultComponentIdGenerator.kt
new file mode 100644
index 0000000..446291b
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/DefaultComponentIdGenerator.kt
@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.compatui.impl
+
+import com.android.wm.shell.compatui.api.CompatUIComponentIdGenerator
+import com.android.wm.shell.compatui.api.CompatUIInfo
+import com.android.wm.shell.compatui.api.CompatUISpec
+
+/**
+ * Default {@link CompatUIComponentIdGenerator} implementation.
+ */
+class DefaultComponentIdGenerator : CompatUIComponentIdGenerator {
+    /**
+     * Simple implementation generating the id from taskId and component name.
+     */
+    override fun generateId(compatUIInfo: CompatUIInfo, spec: CompatUISpec): String =
+        "${compatUIInfo.taskInfo.taskId}-${spec.name}"
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
index 717a414..f22dcce 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
@@ -18,6 +18,7 @@
 
 import static com.android.wm.shell.onehanded.OneHandedController.SUPPORT_ONE_HANDED_MODE;
 
+import android.annotation.NonNull;
 import android.app.ActivityTaskManager;
 import android.content.Context;
 import android.content.pm.PackageManager;
@@ -71,10 +72,13 @@
 import com.android.wm.shell.compatui.CompatUIConfiguration;
 import com.android.wm.shell.compatui.CompatUIController;
 import com.android.wm.shell.compatui.CompatUIShellCommandHandler;
+import com.android.wm.shell.compatui.api.CompatUIComponentIdGenerator;
 import com.android.wm.shell.compatui.api.CompatUIHandler;
 import com.android.wm.shell.compatui.api.CompatUIRepository;
+import com.android.wm.shell.compatui.api.CompatUIState;
 import com.android.wm.shell.compatui.impl.DefaultCompatUIHandler;
 import com.android.wm.shell.compatui.impl.DefaultCompatUIRepository;
+import com.android.wm.shell.compatui.impl.DefaultComponentIdGenerator;
 import com.android.wm.shell.desktopmode.DesktopMode;
 import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
@@ -248,12 +252,15 @@
             Lazy<CompatUIConfiguration> compatUIConfiguration,
             Lazy<CompatUIShellCommandHandler> compatUIShellCommandHandler,
             Lazy<AccessibilityManager> accessibilityManager,
-            CompatUIRepository compatUIRepository) {
+            CompatUIRepository compatUIRepository,
+            @NonNull CompatUIState compatUIState,
+            @NonNull CompatUIComponentIdGenerator componentIdGenerator) {
         if (!context.getResources().getBoolean(R.bool.config_enableCompatUIController)) {
             return Optional.empty();
         }
         if (Flags.appCompatUiFramework()) {
-            return Optional.of(new DefaultCompatUIHandler(compatUIRepository));
+            return Optional.of(new DefaultCompatUIHandler(compatUIRepository, compatUIState,
+                    componentIdGenerator));
         }
         return Optional.of(
                 new CompatUIController(
@@ -274,6 +281,18 @@
 
     @WMSingleton
     @Provides
+    static CompatUIState provideCompatUIState() {
+        return new CompatUIState();
+    }
+
+    @WMSingleton
+    @Provides
+    static CompatUIComponentIdGenerator provideCompatUIComponentIdGenerator() {
+        return new DefaultComponentIdGenerator();
+    }
+
+    @WMSingleton
+    @Provides
     static CompatUIRepository provideCompatUIRepository() {
         return new DefaultCompatUIRepository();
     }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/CompatUIStateUtil.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/CompatUIStateUtil.kt
new file mode 100644
index 0000000..43bd412
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/CompatUIStateUtil.kt
@@ -0,0 +1,51 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.compatui.impl
+
+import com.android.wm.shell.compatui.api.CompatUIComponentState
+import com.android.wm.shell.compatui.api.CompatUISpec
+import com.android.wm.shell.compatui.api.CompatUIState
+import junit.framework.Assert.assertEquals
+import junit.framework.Assert.assertNotNull
+import junit.framework.Assert.assertNull
+
+/**
+ * Asserts no component state exists for the given CompatUISpec
+ */
+internal fun CompatUIState.assertHasNoStateFor(componentId: String) =
+    assertNull(stateForComponent(componentId))
+
+/**
+ * Asserts component state for the given CompatUISpec
+ */
+internal fun CompatUIState.assertHasStateEqualsTo(
+    componentId: String,
+    expected: CompatUIComponentState
+) =
+    assertEquals(stateForComponent(componentId), expected)
+
+/**
+ * Asserts no component exists for the given CompatUISpec
+ */
+internal fun CompatUIState.assertHasNoComponentFor(componentId: String) =
+    assertNull(getUIComponent(componentId))
+
+/**
+ * Asserts component for the given CompatUISpec
+ */
+internal fun CompatUIState.assertHasComponentFor(componentId: String) =
+    assertNotNull(getUIComponent(componentId))
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/DefaultCompatUIHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/DefaultCompatUIHandlerTest.kt
new file mode 100644
index 0000000..8136074
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/DefaultCompatUIHandlerTest.kt
@@ -0,0 +1,196 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.compatui.impl
+
+import android.app.ActivityManager
+import android.testing.AndroidTestingRunner
+import androidx.test.filters.SmallTest
+import com.android.wm.shell.compatui.api.CompatUIComponentState
+import com.android.wm.shell.compatui.api.CompatUIInfo
+import com.android.wm.shell.compatui.api.CompatUIState
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Tests for {@link DefaultCompatUIHandler}.
+ *
+ * Build/Install/Run:
+ *  atest WMShellUnitTests:DefaultCompatUIHandlerTest
+ */
+@RunWith(AndroidTestingRunner::class)
+@SmallTest
+class DefaultCompatUIHandlerTest {
+
+    lateinit var compatUIRepository: FakeCompatUIRepository
+    lateinit var compatUIHandler: DefaultCompatUIHandler
+    lateinit var compatUIState: CompatUIState
+    lateinit var fakeIdGenerator: FakeCompatUIComponentIdGenerator
+
+    @Before
+    fun setUp() {
+        compatUIRepository = FakeCompatUIRepository()
+        compatUIState = CompatUIState()
+        fakeIdGenerator = FakeCompatUIComponentIdGenerator("compId")
+        compatUIHandler = DefaultCompatUIHandler(compatUIRepository, compatUIState,
+            fakeIdGenerator)
+    }
+
+    @Test
+    fun `when creationReturn is false no state is stored`() {
+        // We add a spec to the repository
+        val fakeLifecycle = FakeCompatUILifecyclePredicates(
+            creationReturn = false,
+            removalReturn = false
+        )
+        val fakeCompatUISpec = FakeCompatUISpec("one", fakeLifecycle).getSpec()
+        compatUIRepository.addSpec(fakeCompatUISpec)
+
+        val generatedId = fakeIdGenerator.generatedComponentId
+
+        compatUIHandler.onCompatInfoChanged(testCompatUIInfo())
+
+        fakeIdGenerator.assertGenerateInvocations(1)
+        fakeLifecycle.assertCreationInvocation(1)
+        fakeLifecycle.assertRemovalInvocation(0)
+        fakeLifecycle.assertInitialStateInvocation(0)
+        compatUIState.assertHasNoStateFor(generatedId)
+        compatUIState.assertHasNoComponentFor(generatedId)
+
+        compatUIHandler.onCompatInfoChanged(testCompatUIInfo())
+        fakeLifecycle.assertCreationInvocation(2)
+        fakeLifecycle.assertRemovalInvocation(0)
+        fakeLifecycle.assertInitialStateInvocation(0)
+        compatUIState.assertHasNoStateFor(generatedId)
+        compatUIState.assertHasNoComponentFor(generatedId)
+    }
+
+    @Test
+    fun `when creationReturn is true and no state is created no state is stored`() {
+        // We add a spec to the repository
+        val fakeLifecycle = FakeCompatUILifecyclePredicates(
+            creationReturn = true,
+            removalReturn = false
+        )
+        val fakeCompatUISpec = FakeCompatUISpec("one", fakeLifecycle).getSpec()
+        compatUIRepository.addSpec(fakeCompatUISpec)
+
+        val generatedId = fakeIdGenerator.generatedComponentId
+
+        compatUIHandler.onCompatInfoChanged(testCompatUIInfo())
+
+        fakeLifecycle.assertCreationInvocation(1)
+        fakeLifecycle.assertRemovalInvocation(0)
+        fakeLifecycle.assertInitialStateInvocation(1)
+        compatUIState.assertHasNoStateFor(generatedId)
+        compatUIState.assertHasComponentFor(generatedId)
+
+        compatUIHandler.onCompatInfoChanged(testCompatUIInfo())
+
+        fakeLifecycle.assertCreationInvocation(1)
+        fakeLifecycle.assertRemovalInvocation(1)
+        fakeLifecycle.assertInitialStateInvocation(1)
+        compatUIState.assertHasNoStateFor(generatedId)
+        compatUIState.assertHasComponentFor(generatedId)
+    }
+
+    @Test
+    fun `when creationReturn is true and state is created state is stored`() {
+        val fakeComponentState = object : CompatUIComponentState {}
+        // We add a spec to the repository
+        val fakeLifecycle = FakeCompatUILifecyclePredicates(
+            creationReturn = true,
+            removalReturn = false,
+            initialState = { _, _ -> fakeComponentState }
+        )
+        val fakeCompatUISpec = FakeCompatUISpec("one", fakeLifecycle).getSpec()
+        compatUIRepository.addSpec(fakeCompatUISpec)
+
+        val generatedId = fakeIdGenerator.generatedComponentId
+
+        compatUIHandler.onCompatInfoChanged(testCompatUIInfo())
+
+        fakeLifecycle.assertCreationInvocation(1)
+        fakeLifecycle.assertRemovalInvocation(0)
+        fakeLifecycle.assertInitialStateInvocation(1)
+        compatUIState.assertHasStateEqualsTo(generatedId, fakeComponentState)
+        compatUIState.assertHasComponentFor(generatedId)
+
+        compatUIHandler.onCompatInfoChanged(testCompatUIInfo())
+
+        fakeLifecycle.assertCreationInvocation(1)
+        fakeLifecycle.assertRemovalInvocation(1)
+        fakeLifecycle.assertInitialStateInvocation(1)
+        compatUIState.assertHasStateEqualsTo(generatedId, fakeComponentState)
+        compatUIState.assertHasComponentFor(generatedId)
+    }
+
+    @Test
+    fun `when lifecycle is complete and state is created state is stored and removed`() {
+        val fakeComponentState = object : CompatUIComponentState {}
+        // We add a spec to the repository
+        val fakeLifecycle = FakeCompatUILifecyclePredicates(
+            creationReturn = true,
+            removalReturn = true,
+            initialState = { _, _ -> fakeComponentState }
+        )
+        val fakeCompatUISpec = FakeCompatUISpec("one", fakeLifecycle).getSpec()
+        compatUIRepository.addSpec(fakeCompatUISpec)
+
+        val generatedId = fakeIdGenerator.generatedComponentId
+
+        compatUIHandler.onCompatInfoChanged(testCompatUIInfo())
+
+        fakeLifecycle.assertCreationInvocation(1)
+        fakeLifecycle.assertRemovalInvocation(0)
+        fakeLifecycle.assertInitialStateInvocation(1)
+        compatUIState.assertHasStateEqualsTo(generatedId, fakeComponentState)
+        compatUIState.assertHasComponentFor(generatedId)
+
+        compatUIHandler.onCompatInfoChanged(testCompatUIInfo())
+
+        fakeLifecycle.assertCreationInvocation(1)
+        fakeLifecycle.assertRemovalInvocation(1)
+        fakeLifecycle.assertInitialStateInvocation(1)
+        compatUIState.assertHasNoStateFor(generatedId)
+        compatUIState.assertHasNoComponentFor(generatedId)
+    }
+
+    @Test
+    fun `idGenerator is invoked every time a component is created`() {
+        // We add a spec to the repository
+        val fakeLifecycle = FakeCompatUILifecyclePredicates(
+            creationReturn = true,
+            removalReturn = true,
+        )
+        val fakeCompatUISpec = FakeCompatUISpec("one", fakeLifecycle).getSpec()
+        compatUIRepository.addSpec(fakeCompatUISpec)
+        // Component creation
+        fakeIdGenerator.assertGenerateInvocations(0)
+        compatUIHandler.onCompatInfoChanged(testCompatUIInfo())
+        fakeIdGenerator.assertGenerateInvocations(1)
+
+        compatUIHandler.onCompatInfoChanged(testCompatUIInfo())
+        fakeIdGenerator.assertGenerateInvocations(2)
+    }
+
+    private fun testCompatUIInfo(): CompatUIInfo {
+        val taskInfo = ActivityManager.RunningTaskInfo()
+        taskInfo.taskId = 1
+        return CompatUIInfo(taskInfo, null)
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/DefaultCompatUIRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/DefaultCompatUIRepositoryTest.kt
index 1a86cfd..e35acb2 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/DefaultCompatUIRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/DefaultCompatUIRepositoryTest.kt
@@ -19,6 +19,7 @@
 import android.platform.test.flag.junit.DeviceFlagsValueProvider
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
+import com.android.wm.shell.compatui.api.CompatUILifecyclePredicates
 import com.android.wm.shell.compatui.api.CompatUIRepository
 import com.android.wm.shell.compatui.api.CompatUISpec
 import org.junit.Assert.assertEquals
@@ -50,16 +51,16 @@
 
     @Test(expected = IllegalStateException::class)
     fun `addSpec throws exception with specs with duplicate id`() {
-        repository.addSpec(CompatUISpec("one"))
-        repository.addSpec(CompatUISpec("one"))
+        repository.addSpec(specById("one"))
+        repository.addSpec(specById("one"))
     }
 
     @Test
     fun `iterateOn invokes the consumer`() {
         with(repository) {
-            addSpec(CompatUISpec("one"))
-            addSpec(CompatUISpec("two"))
-            addSpec(CompatUISpec("three"))
+            addSpec(specById("one"))
+            addSpec(specById("two"))
+            addSpec(specById("three"))
             val consumer = object : (CompatUISpec) -> Unit {
                 var acc = ""
                 override fun invoke(spec: CompatUISpec) {
@@ -74,9 +75,9 @@
     @Test
     fun `findSpec returns existing specs`() {
         with(repository) {
-            val one = CompatUISpec("one")
-            val two = CompatUISpec("two")
-            val three = CompatUISpec("three")
+            val one = specById("one")
+            val two = specById("two")
+            val three = specById("three")
             addSpec(one)
             addSpec(two)
             addSpec(three)
@@ -86,4 +87,10 @@
             assertNull(findSpec("abc"))
         }
     }
+
+    private fun specById(name: String): CompatUISpec =
+        CompatUISpec(name = name, lifecycle = CompatUILifecyclePredicates(
+            creationPredicate = { _, _ -> true },
+            removalPredicate = { _, _, _ -> true }
+        ))
 }
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/FakeCompatUIComponentIdGenerator.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/FakeCompatUIComponentIdGenerator.kt
new file mode 100644
index 0000000..bc743ed
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/FakeCompatUIComponentIdGenerator.kt
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.compatui.impl
+
+import com.android.wm.shell.compatui.api.CompatUIComponentIdGenerator
+import com.android.wm.shell.compatui.api.CompatUIInfo
+import com.android.wm.shell.compatui.api.CompatUISpec
+import junit.framework.Assert.assertEquals
+
+/**
+ * A Fake {@link CompatUIComponentIdGenerator} implementation.
+ */
+class FakeCompatUIComponentIdGenerator(var generatedComponentId: String = "compId") :
+    CompatUIComponentIdGenerator {
+
+    var generateInvocations = 0
+
+    override fun generateId(compatUIInfo: CompatUIInfo, spec: CompatUISpec): String {
+        generateInvocations++
+        return generatedComponentId
+    }
+
+    fun resetInvocations() {
+        generateInvocations = 0
+    }
+
+    fun assertGenerateInvocations(expected: Int) =
+        assertEquals(expected, generateInvocations)
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/FakeCompatUILifecyclePredicates.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/FakeCompatUILifecyclePredicates.kt
new file mode 100644
index 0000000..bbaa2db
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/FakeCompatUILifecyclePredicates.kt
@@ -0,0 +1,69 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.compatui.impl
+
+import com.android.wm.shell.compatui.api.CompatUIComponentState
+import com.android.wm.shell.compatui.api.CompatUIInfo
+import com.android.wm.shell.compatui.api.CompatUILifecyclePredicates
+import com.android.wm.shell.compatui.api.CompatUISharedState
+import junit.framework.Assert.assertEquals
+
+/**
+ * Fake class for {@link CompatUILifecycle}
+ */
+class FakeCompatUILifecyclePredicates(
+    private val creationReturn: Boolean,
+    private val removalReturn: Boolean,
+    private val initialState: (
+        CompatUIInfo,
+        CompatUISharedState
+    ) -> CompatUIComponentState? = { _, _ -> null }
+) {
+    var creationInvocation = 0
+    var removalInvocation = 0
+    var initialStateInvocation = 0
+    var lastCreationCompatUIInfo: CompatUIInfo? = null
+    var lastCreationSharedState: CompatUISharedState? = null
+    var lastRemovalCompatUIInfo: CompatUIInfo? = null
+    var lastRemovalSharedState: CompatUISharedState? = null
+    var lastRemovalCompState: CompatUIComponentState? = null
+    fun getLifecycle() = CompatUILifecyclePredicates(
+        creationPredicate = { uiInfo, sharedState ->
+            lastCreationCompatUIInfo = uiInfo
+            lastCreationSharedState = sharedState
+            creationInvocation++
+            creationReturn
+        },
+        removalPredicate = { uiInfo, sharedState, compState ->
+            lastRemovalCompatUIInfo = uiInfo
+            lastRemovalSharedState = sharedState
+            lastRemovalCompState = compState
+            removalInvocation++
+            removalReturn
+        },
+        stateBuilder = { a, b -> initialStateInvocation++; initialState(a, b) }
+    )
+
+    fun assertCreationInvocation(expected: Int) =
+        assertEquals(expected, creationInvocation)
+
+    fun assertRemovalInvocation(expected: Int) =
+        assertEquals(expected, removalInvocation)
+
+    fun assertInitialStateInvocation(expected: Int) =
+        assertEquals(expected, initialStateInvocation)
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/FakeCompatUISpec.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/FakeCompatUISpec.kt
new file mode 100644
index 0000000..1ecd52e
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/impl/FakeCompatUISpec.kt
@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.compatui.impl
+
+import com.android.wm.shell.compatui.api.CompatUISpec
+
+/**
+ * Fake implementation for {@link ompatUISpec}
+ */
+class FakeCompatUISpec(
+    val name: String,
+    val lifecycle: FakeCompatUILifecyclePredicates
+) {
+    fun getSpec(): CompatUISpec = CompatUISpec(
+        name = name,
+        lifecycle = lifecycle.getLifecycle()
+    )
+}