From ee07f627c947a07533b1f5d0d5e559a17e934e99 Mon Sep 17 00:00:00 2001
From: Narayan Kamath <narayan@google.com>
Date: Mon, 8 Jan 2018 12:20:28 +0000
Subject: [PATCH] WorkSource: Fix WorkSource#remove for chained worksources.

- It doesn't make sense to clear the list of WorkChains if we're
  remove is called with a WorkSource that has no chains.
- The early return for mNum <= 0 is faulty because it incorrectly
  returns early for WorkSources that have workChains but no flat
  UIDs.

Test: WorkSourceTest
Bug: 62390666
Change-Id: I4c0e69bdd7c114b41329aa329d1c1d08a8cb9b59
---
 core/java/android/os/WorkSource.java          | 13 ++++--------
 .../src/android/os/WorkSourceTest.java        | 20 +++++++++++++++++++
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/core/java/android/os/WorkSource.java b/core/java/android/os/WorkSource.java
index 401b4a36a7438..bb043b0109d5f 100644
--- a/core/java/android/os/WorkSource.java
+++ b/core/java/android/os/WorkSource.java
@@ -407,11 +407,11 @@ public class WorkSource implements Parcelable {
     }
 
     public boolean remove(WorkSource other) {
-        if (mNum <= 0 || other.mNum <= 0) {
+        if (isEmpty() || other.isEmpty()) {
             return false;
         }
 
-        boolean uidRemoved = false;
+        boolean uidRemoved;
         if (mNames == null && other.mNames == null) {
             uidRemoved = removeUids(other);
         } else {
@@ -427,13 +427,8 @@ public class WorkSource implements Parcelable {
         }
 
         boolean chainRemoved = false;
-        if (other.mChains != null) {
-            if (mChains != null) {
-                chainRemoved = mChains.removeAll(other.mChains);
-            }
-        } else if (mChains != null) {
-            mChains.clear();
-            chainRemoved = true;
+        if (other.mChains != null && mChains != null) {
+            chainRemoved = mChains.removeAll(other.mChains);
         }
 
         return uidRemoved || chainRemoved;
diff --git a/core/tests/coretests/src/android/os/WorkSourceTest.java b/core/tests/coretests/src/android/os/WorkSourceTest.java
index 90b457561180e..566ac4daf9502 100644
--- a/core/tests/coretests/src/android/os/WorkSourceTest.java
+++ b/core/tests/coretests/src/android/os/WorkSourceTest.java
@@ -331,4 +331,24 @@ public class WorkSourceTest extends TestCase {
         wc.addNode(200, "tag2");
         assertEquals(100, wc.getAttributionUid());
     }
+
+    public void testRemove_fromChainedWorkSource() {
+        WorkSource ws1 = new WorkSource();
+        ws1.createWorkChain().addNode(50, "foo");
+        ws1.createWorkChain().addNode(75, "bar");
+        ws1.add(100);
+
+        WorkSource ws2 = new WorkSource();
+        ws2.add(100);
+
+        assertTrue(ws1.remove(ws2));
+        assertEquals(2, ws1.getWorkChains().size());
+        assertEquals(50, ws1.getWorkChains().get(0).getAttributionUid());
+        assertEquals(75, ws1.getWorkChains().get(1).getAttributionUid());
+
+        ws2.createWorkChain().addNode(50, "foo");
+        assertTrue(ws1.remove(ws2));
+        assertEquals(1, ws1.getWorkChains().size());
+        assertEquals(75, ws1.getWorkChains().get(0).getAttributionUid());
+    }
 }