Merge "[4/n] Create first state and handle comp lifecycle" into main
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()
+    )
+}