Have getFragments() actually return all fragments

Change-Id: Ib255a4087122906ac2272d802fa5fd9402ae8a90

3 files changed, 16 insertions, 8 deletions

diff --git a/src/main/java/com/xtremelabs/robolectric/shadows/ShadowFragmentActivity.java b/src/main/java/com/xtremelabs/robolectric/shadows/ShadowFragmentActivity.java
index 3071162ed..98172b7f7 100644
--- a/src/main/java/com/xtremelabs/robolectric/shadows/ShadowFragmentActivity.java
+++ b/src/main/java/com/xtremelabs/robolectric/shadows/ShadowFragmentActivity.java
@@ -55,14 +55,14 @@ public class ShadowFragmentActivity extends ShadowActivity {
 
     @Implementation
     public void onStart() {
-        for (Fragment fragment : fragmentManager.getFragments().values()) {
+        for (Fragment fragment : fragmentManager.getFragments()) {
             fragmentManager.startFragment(fragment);
         }
     }
 
     @Implementation
     public void onPause() {
-        for(Fragment fragment : fragmentManager.getFragments().values()) {
+        for(Fragment fragment : fragmentManager.getFragments()) {
             fragment.onPause();
         }
     }
@@ -78,7 +78,7 @@ public class ShadowFragmentActivity extends ShadowActivity {
         // terrible looking hack.  I am very sorry.
         List<SerializedFragmentState> fragmentStates = new ArrayList<SerializedFragmentState>();
 
-        for (Map.Entry<Integer, Fragment> entry : fragmentManager.getFragments().entrySet()) {
+        for (Map.Entry<Integer, Fragment> entry : fragmentManager.getFragmentsById().entrySet()) {
             Fragment fragment = entry.getValue();
             fragment.onSaveInstanceState(outState);
             fragmentStates.add(new SerializedFragmentState(entry.getKey(), fragment));
@@ -95,7 +95,7 @@ public class ShadowFragmentActivity extends ShadowActivity {
             return focusedView;
         }
 
-        for (Fragment fragment : fragmentManager.getFragments().values()) {
+        for (Fragment fragment : fragmentManager.getFragments()) {
             View view = shadowOf(fragment).view;
             if (view != null) {
                 focusedView = view.findFocus();
@@ -111,7 +111,7 @@ public class ShadowFragmentActivity extends ShadowActivity {
     public void clearFocus() {
         super.clearFocus();
 
-        for (Fragment fragment : fragmentManager.getFragments().values()) {
+        for (Fragment fragment : fragmentManager.getFragments()) {
             View view = shadowOf(fragment).view;
             if (view != null) {
                 view.clearFocus();
diff --git a/src/main/java/com/xtremelabs/robolectric/tester/android/util/TestFragmentManager.java b/src/main/java/com/xtremelabs/robolectric/tester/android/util/TestFragmentManager.java
index 3aa2ece10..115cde4a6 100644
--- a/src/main/java/com/xtremelabs/robolectric/tester/android/util/TestFragmentManager.java
+++ b/src/main/java/com/xtremelabs/robolectric/tester/android/util/TestFragmentManager.java
@@ -15,6 +15,7 @@ import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 
@@ -175,10 +176,17 @@ public class TestFragmentManager extends FragmentManager {
         fragment.onStart();
     }
 
-    public HashMap<Integer, Fragment> getFragments() {
+    public HashMap<Integer, Fragment> getFragmentsById() {
         return new HashMap<Integer, Fragment>(fragmentsById);
     }
 
+    public HashSet<Fragment> getFragments() {
+        HashSet<Fragment> fragments = new HashSet<Fragment>();
+        fragments.addAll(fragmentsById.values());
+        fragments.addAll(fragmentsByTag.values());
+        return fragments;
+    }
+
     public List<TestFragmentTransaction> getCommittedTransactions() {
         return transactions;
     }
diff --git a/src/test/java/com/xtremelabs/robolectric/shadows/FragmentActivityTest.java b/src/test/java/com/xtremelabs/robolectric/shadows/FragmentActivityTest.java
index d03deed34..557757877 100644
--- a/src/test/java/com/xtremelabs/robolectric/shadows/FragmentActivityTest.java
+++ b/src/test/java/com/xtremelabs/robolectric/shadows/FragmentActivityTest.java
@@ -116,7 +116,7 @@ public class FragmentActivityTest {
         TestFragmentManager fragmentManager = (TestFragmentManager) activity.getSupportFragmentManager();
         assertEquals(2, fragmentManager.getFragments().size());
 
-        TestFragment restoredFrag = (TestFragment) fragmentManager.getFragments().get(containerId);
+        TestFragment restoredFrag = (TestFragment) fragmentManager.getFragmentsById().get(containerId);
         assertEquals(restoredFrag.getId(), dynamicFrag.getId());
         assertEquals(restoredFrag.getTag(), dynamicFrag.getTag());
         assertEquals(bundle, shadowOf(restoredFrag).getSavedInstanceState());
@@ -138,7 +138,7 @@ public class FragmentActivityTest {
         shadowOf(activity).onStart();
         TestFragmentManager fragmentManager = (TestFragmentManager) activity.getSupportFragmentManager();
         assertEquals(2, fragmentManager.getFragments().size());
-        TestFragment restoredFrag = (TestFragment) fragmentManager.getFragments().get(containerId);
+        TestFragment restoredFrag = (TestFragment) fragmentManager.getFragmentsById().get(containerId);
 
         assertEquals(restoredFrag.onCreateViewInflater, activity.getLayoutInflater());
         assertNotNull(restoredFrag.getView());