5 files changed, 84 insertions, 21 deletions

diff --git a/updater_sample/tests/res/raw/update_config_stream_001.json b/updater_sample/tests/res/raw/update_config_001_stream.json
index be51b7c95..be51b7c95 100644
--- a/updater_sample/tests/res/raw/update_config_stream_001.json
+++ b/updater_sample/tests/res/raw/update_config_001_stream.json
diff --git a/updater_sample/tests/res/raw/update_config_stream_002.json b/updater_sample/tests/res/raw/update_config_002_stream.json
index 5d7874cdb..5d7874cdb 100644
--- a/updater_sample/tests/res/raw/update_config_stream_002.json
+++ b/updater_sample/tests/res/raw/update_config_002_stream.json
diff --git a/updater_sample/tests/res/raw/update_config_003_nonstream.json b/updater_sample/tests/res/raw/update_config_003_nonstream.json
new file mode 100644
index 000000000..4175c35ea
--- /dev/null
+++ b/updater_sample/tests/res/raw/update_config_003_nonstream.json
@@ -0,0 +1,9 @@
+{
+    "__": "*** Generated using tools/gen_update_config.py ***",
+    "ab_config": {
+        "force_switch_slot": false
+    },
+    "ab_install_type": "NON_STREAMING",
+    "name": "S ota_002_package",
+    "url": "file:///data/sample-ota-packages/ota_003_package.zip"
+}
\ No newline at end of file
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
index 000f5663b..1cbd8601e 100644
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
@@ -60,7 +60,7 @@ public class UpdateConfigTest {
     public void setUp() throws Exception {
         mContext = InstrumentationRegistry.getContext();
         mTargetContext = InstrumentationRegistry.getTargetContext();
-        mJsonStreaming001 = readResource(R.raw.update_config_stream_001);
+        mJsonStreaming001 = readResource(R.raw.update_config_001_stream);
     }
 
     @Test
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
index 0657a5eb6..e05ad290c 100644
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
@@ -18,19 +18,22 @@ package com.example.android.systemupdatersample;
 
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.content.Context;
 import android.os.UpdateEngine;
 import android.os.UpdateEngineCallback;
+import android.support.test.InstrumentationRegistry;
 import android.support.test.filters.SmallTest;
 import android.support.test.runner.AndroidJUnit4;
 
+import com.example.android.systemupdatersample.tests.R;
 import com.example.android.systemupdatersample.util.PayloadSpecs;
+import com.google.common.collect.ImmutableList;
+import com.google.common.io.CharStreams;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -40,7 +43,9 @@ import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
-import java.util.function.IntConsumer;
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStreamReader;
 
 /**
  * Tests for {@link UpdateManager}
@@ -56,37 +61,86 @@ public class UpdateManagerTest {
     private UpdateEngine mUpdateEngine;
     @Mock
     private PayloadSpecs mPayloadSpecs;
-    private UpdateManager mUpdateManager;
+    private UpdateManager mSubject;
+    private Context mContext;
+    private UpdateConfig mNonStreamingUpdate003;
 
     @Before
-    public void setUp() {
-        mUpdateManager = new UpdateManager(mUpdateEngine, mPayloadSpecs);
+    public void setUp() throws Exception {
+        mContext = InstrumentationRegistry.getContext();
+        mSubject = new UpdateManager(mUpdateEngine, mPayloadSpecs);
+        mNonStreamingUpdate003 =
+                UpdateConfig.fromJson(readResource(R.raw.update_config_003_nonstream));
     }
 
     @Test
-    public void storesProgressThenInvokesCallbacks() {
-        IntConsumer statusUpdateCallback = mock(IntConsumer.class);
-
-        // When UpdateManager is bound to update_engine, it passes
-        // UpdateManager.UpdateEngineCallbackImpl as a callback to update_engine.
+    public void applyUpdate_appliesPayloadToUpdateEngine() throws Exception {
+        PayloadSpec payload = buildMockPayloadSpec();
+        when(mPayloadSpecs.forNonStreaming(any(File.class))).thenReturn(payload);
         when(mUpdateEngine.bind(any(UpdateEngineCallback.class))).thenAnswer(answer -> {
+            // When UpdateManager is bound to update_engine, it passes
+            // UpdateEngineCallback as a callback to update_engine.
             UpdateEngineCallback callback = answer.getArgument(0);
-            callback.onStatusUpdate(/*engineStatus*/ 4, /*engineProgress*/ 0.2f);
+            callback.onStatusUpdate(
+                    UpdateEngine.UpdateStatusConstants.IDLE,
+                    /*engineProgress*/ 0.0f);
             return null;
         });
 
-        mUpdateManager.setOnEngineStatusUpdateCallback(statusUpdateCallback);
+        mSubject.bind();
+        mSubject.applyUpdate(null, mNonStreamingUpdate003);
+
+        verify(mUpdateEngine).applyPayload(
+                "file://blah",
+                120,
+                340,
+                new String[] {
+                        "SWITCH_SLOT_ON_REBOOT=0" // ab_config.force_switch_slot = false
+                });
+    }
 
-        // Making sure that manager.getProgress() returns correct progress
-        // in "onEngineStatusUpdate" callback.
-        doAnswer(answer -> {
-            assertEquals(0.2f, mUpdateManager.getProgress(), 1E-5);
+    @Test
+    public void stateIsRunningAndEngineStatusIsIdle_reApplyLastUpdate() throws Exception {
+        PayloadSpec payload = buildMockPayloadSpec();
+        when(mPayloadSpecs.forNonStreaming(any(File.class))).thenReturn(payload);
+        when(mUpdateEngine.bind(any(UpdateEngineCallback.class))).thenAnswer(answer -> {
+            // When UpdateManager is bound to update_engine, it passes
+            // UpdateEngineCallback as a callback to update_engine.
+            UpdateEngineCallback callback = answer.getArgument(0);
+            callback.onStatusUpdate(
+                    UpdateEngine.UpdateStatusConstants.IDLE,
+                    /*engineProgress*/ 0.0f);
             return null;
-        }).when(statusUpdateCallback).accept(anyInt());
+        });
 
-        mUpdateManager.bind();
+        mSubject.bind();
+        mSubject.applyUpdate(null, mNonStreamingUpdate003);
+        mSubject.unbind();
+        mSubject.bind(); // re-bind - now it should re-apply last update
+
+        assertEquals(mSubject.getUpdaterState(), UpdaterState.RUNNING);
+        // it should be called 2 times
+        verify(mUpdateEngine, times(2)).applyPayload(
+                "file://blah",
+                120,
+                340,
+                new String[] {
+                        "SWITCH_SLOT_ON_REBOOT=0" // ab_config.force_switch_slot = false
+                });
+    }
+
+    private PayloadSpec buildMockPayloadSpec() {
+        PayloadSpec payload = mock(PayloadSpec.class);
+        when(payload.getUrl()).thenReturn("file://blah");
+        when(payload.getOffset()).thenReturn(120L);
+        when(payload.getSize()).thenReturn(340L);
+        when(payload.getProperties()).thenReturn(ImmutableList.of());
+        return payload;
+    }
 
-        verify(statusUpdateCallback, times(1)).accept(4);
+    private String readResource(int id) throws IOException {
+        return CharStreams.toString(new InputStreamReader(
+                mContext.getResources().openRawResource(id)));
     }
 
 }