Apply badge state restoration on the view itself

.gitignore

@@ -38,4 +38,12 @@ captures/
 .idea/libraries
 
 # Keystore files
-*.jks
+*.jks
+
+
+# IntelliJ
+*.iml
+*.ipr
+*.iws
+.idea/
+out/

app/src/main/java/com/example/bottombar/sample/BadgeActivity.java

@@ -17,6 +17,7 @@
  */
 public class BadgeActivity extends AppCompatActivity {
     private TextView messageView;
+    private BottomBar bottomBar;
 
     @Override
     protected void onCreate(@Nullable Bundle savedInstanceState) {
@@ -25,7 +26,7 @@ protected void onCreate(@Nullable Bundle savedInstanceState) {
 
         messageView = (TextView) findViewById(R.id.messageView);
 
-        final BottomBar bottomBar = (BottomBar) findViewById(R.id.bottomBar);
+        bottomBar = (BottomBar) findViewById(R.id.bottomBar);
         bottomBar.setOnTabSelectListener(new OnTabSelectListener() {
             @Override
             public void onTabSelected(@IdRes int tabId) {
@@ -39,7 +40,11 @@ public void onTabReSelected(@IdRes int tabId) {
                 Toast.makeText(getApplicationContext(), TabMessage.get(tabId, true), Toast.LENGTH_LONG).show();
             }
         });
+    }
 
+    @Override
+    protected void onResume() {
+        super.onResume();
         BottomBarTab nearby = bottomBar.getTabWithId(R.id.tab_nearby);
         nearby.setBadgeCount(5);
     }

bottom-bar/src/androidTest/java/com/roughike/bottombar/BadgeTest.java

@@ -1,12 +1,11 @@
 package com.roughike.bottombar;
 
-import android.os.Bundle;
+import android.os.Parcelable;
 import android.support.test.InstrumentationRegistry;
 import android.support.test.annotation.UiThreadTest;
 import android.support.test.filters.LargeTest;
 import android.support.test.runner.AndroidJUnit4;
 
-import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -75,13 +74,17 @@ public void whenBadgeStateRestored_CountPersists() {
         nearby.setBadgeCount(1);
         assertEquals(1, nearby.badge.getCount());
 
+        nearby.badge.setCount(2);
+        Parcelable stateGood = nearby.badge.onSaveInstanceState();
+        nearby.badge.onRestoreInstanceState(stateGood);
+        assertEquals(2, nearby.badge.getCount());
 
-        int tabIndex = nearby.getIndexInTabContainer();
-        Bundle savedInstanceState = new Bundle();
-        savedInstanceState.putInt(BottomBarBadge.STATE_COUNT + tabIndex, 2);
-        nearby.badge.restoreState(savedInstanceState, tabIndex);
+        nearby.badge.setCount(4);
+        BottomBarBadge barBadge = new BottomBarBadge(InstrumentationRegistry.getContext());
+        Parcelable stateBad = barBadge.onSaveInstanceState();
+        nearby.badge.onRestoreInstanceState(stateBad);
 
-        assertEquals(2, nearby.badge.getCount());
+        assertNotEquals(4, nearby.badge.getCount());
     }
 
     @Test

bottom-bar/src/androidTest/java/com/roughike/bottombar/BottomBarTabTest.java

@@ -1,6 +1,5 @@
 package com.roughike.bottombar;
 
-import android.os.Bundle;
 import android.support.test.InstrumentationRegistry;
 import android.support.test.runner.AndroidJUnit4;
 import android.widget.FrameLayout;
@@ -10,6 +9,8 @@
 import org.junit.runner.RunWith;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 
 /**
  * Created by iiro on 22.8.2016.
@@ -51,13 +52,20 @@ public void testSavedStateWithBadge_StaysIntact() {
         tab.setIndexInContainer(69);
         assertEquals(69, tab.getIndexInTabContainer());
 
-        Bundle savedState = (Bundle) tab.onSaveInstanceState();
-        assertEquals(5, savedState.getInt(BottomBarBadge.STATE_COUNT + 69));
-
         tab.setBadgeCount(9);
         assertEquals(9, tab.badge.getCount());
+    }
 
-        tab.onRestoreInstanceState(savedState);
-        assertEquals(5, tab.badge.getCount());
+    @Test
+    public void testBadgeVisibility_BadgeNull() {
+        tab.setBadgeCount(0);
+        assertNull(tab.badge);
+    }
+
+    @Test
+    public void testBadgeVisibility_BadgeNotNull() {
+        tab.setBadgeCount(0);
+        tab.setBadgeCount(5);
+        assertNotNull(tab.badge);
     }
 }

bottom-bar/src/main/java/com/roughike/bottombar/BottomBarBadge.java

@@ -4,8 +4,8 @@
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.ShapeDrawable;
 import android.os.Build;
-import android.os.Bundle;
-import android.support.annotation.VisibleForTesting;
+import android.os.Parcel;
+import android.os.Parcelable;
 import android.support.v4.view.ViewCompat;
 import android.support.v7.widget.AppCompatImageView;
 import android.view.Gravity;
@@ -31,8 +31,6 @@
  * limitations under the License.
  */
 class BottomBarBadge extends TextView {
-    @VisibleForTesting
-    static final String STATE_COUNT = "STATE_BADGE_COUNT_FOR_TAB_";
 
     private int count;
     private boolean isVisible = false;
@@ -176,13 +174,53 @@ private void setBackgroundCompat(Drawable background) {
         }
     }
 
-    Bundle saveState(int tabIndex) {
-        Bundle state = new Bundle();
-        state.putInt(STATE_COUNT + tabIndex, count);
-        return state;
+    @Override
+    public Parcelable onSaveInstanceState() {
+        Parcelable superState = super.onSaveInstanceState();
+        SavedState ss = new SavedState(superState);
+        ss.savedCount = this.count;
+        return ss;
     }
 
-    void restoreState(Bundle bundle, int tabIndex) {
-        setCount(bundle.getInt(STATE_COUNT + tabIndex, count));
+    @Override
+    public void onRestoreInstanceState(Parcelable state) {
+        if (!(state instanceof SavedState)) {
+            super.onRestoreInstanceState(state);
+            return;
+        }
+        SavedState ss = (SavedState) state;
+        super.onRestoreInstanceState(ss.getSuperState());
+        this.count = ss.savedCount;
+        setCount(count);
+    }
+
+    public static class SavedState extends BaseSavedState {
+        int savedCount;
+
+        SavedState(Parcelable superState) {
+            super(superState);
+        }
+
+        private SavedState(Parcel in) {
+            super(in);
+            this.savedCount = in.readInt();
+        }
+
+        @Override
+        public void writeToParcel(Parcel out, int flags) {
+            super.writeToParcel(out, flags);
+            out.writeInt(this.savedCount);
+        }
+
+        public static final Parcelable.Creator<SavedState> CREATOR =
+                new Parcelable.Creator<SavedState>() {
+                    public SavedState createFromParcel(Parcel in) {
+                        return new SavedState(in);
+                    }
+
+                    public SavedState[] newArray(int size) {
+                        return new SavedState[size];
+                    }
+                };
     }
 }

bottom-bar/src/main/java/com/roughike/bottombar/BottomBarTab.java

@@ -7,8 +7,6 @@
 import android.content.Context;
 import android.graphics.Typeface;
 import android.os.Build;
-import android.os.Bundle;
-import android.os.Parcelable;
 import android.support.annotation.ColorInt;
 import android.support.annotation.VisibleForTesting;
 import android.support.v4.view.ViewCompat;
@@ -533,28 +531,6 @@ private void setTitleScale(float scale) {
         ViewCompat.setScaleY(titleView, scale);
     }
 
-    @Override
-    public Parcelable onSaveInstanceState() {
-        if (badge != null) {
-            Bundle bundle = badge.saveState(indexInContainer);
-            bundle.putParcelable("superstate", super.onSaveInstanceState());
-            return bundle;
-        }
-
-        return super.onSaveInstanceState();
-    }
-
-    @Override
-    public void onRestoreInstanceState(Parcelable state) {
-        if (badge != null && state instanceof Bundle) {
-            Bundle bundle = (Bundle) state;
-            badge.restoreState(bundle, indexInContainer);
-
-            state = bundle.getParcelable("superstate");
-        }
-        super.onRestoreInstanceState(state);
-    }
-
     public static class Config {
         private final float inActiveTabAlpha;
         private final float activeTabAlpha;