Merge "Complete the proto events if it's missing due to interruption"
diff --git a/src/com/android/tradefed/result/proto/ProtoResultParser.java b/src/com/android/tradefed/result/proto/ProtoResultParser.java
index 5211aad..0f775f9 100644
--- a/src/com/android/tradefed/result/proto/ProtoResultParser.java
+++ b/src/com/android/tradefed/result/proto/ProtoResultParser.java
@@ -37,6 +37,7 @@
 import com.android.tradefed.result.proto.TestRecordProto.ChildReference;
 import com.android.tradefed.result.proto.TestRecordProto.DebugInfo;
 import com.android.tradefed.result.proto.TestRecordProto.DebugInfoContext;
+import com.android.tradefed.result.proto.TestRecordProto.FailureStatus;
 import com.android.tradefed.result.proto.TestRecordProto.TestRecord;
 import com.android.tradefed.testtype.suite.ModuleDefinition;
 import com.android.tradefed.util.MultiMap;
@@ -77,6 +78,8 @@
     private boolean mInvocationStarted = false;
     private boolean mInvocationEnded = false;
     private boolean mFirstModule = true;
+    /** Track the name of the module in progress. */
+    private String mModuleInProgress = null;
 
     /** Ctor. */
     public ProtoResultParser(
@@ -189,6 +192,26 @@
         return mInvocationEnded;
     }
 
+    /** Returns the id of the module in progress. Returns null if none in progress. */
+    public String getModuleInProgress() {
+        return mModuleInProgress;
+    }
+
+    /** If needed to ensure consistent reporting, complete the events of the module. */
+    public void completeModuleEvents() {
+        if (getModuleInProgress() == null) {
+            return;
+        }
+        mListener.testRunStarted(getModuleInProgress(), 0);
+        FailureDescription failure =
+                FailureDescription.create(
+                        "Module was interrupted after starting, results are incomplete.",
+                        FailureStatus.INFRA_FAILURE);
+        mListener.testRunFailed(failure);
+        mListener.testRunEnded(0L, new HashMap<String, Metric>());
+        mListener.testModuleEnded();
+    }
+
     private void evalChildrenProto(List<ChildReference> children, boolean isInRun) {
         for (ChildReference child : children) {
             TestRecord childProto = child.getInlineTestRecord();
@@ -328,12 +351,13 @@
                     InvocationContext.fromProto(anyDescription.unpack(Context.class));
             String message = "Test module started proto";
             if (moduleContext.getAttributes().containsKey(ModuleDefinition.MODULE_ID)) {
-                message +=
-                        (": "
-                                + moduleContext
-                                        .getAttributes()
-                                        .getUniqueMap()
-                                        .get(ModuleDefinition.MODULE_ID));
+                String moduleId =
+                        moduleContext
+                                .getAttributes()
+                                .getUniqueMap()
+                                .get(ModuleDefinition.MODULE_ID);
+                message += (": " + moduleId);
+                mModuleInProgress = moduleId;
             }
             log(message);
             mListener.testModuleStarted(moduleContext);
@@ -351,6 +375,7 @@
         handleLogs(moduleProto);
         log("Test module ended proto");
         mListener.testModuleEnded();
+        mModuleInProgress = null;
     }
 
     /** Handles the test run level of the invocation. */
diff --git a/src/com/android/tradefed/result/proto/StreamProtoReceiver.java b/src/com/android/tradefed/result/proto/StreamProtoReceiver.java
index adef18b..df6df56 100644
--- a/src/com/android/tradefed/result/proto/StreamProtoReceiver.java
+++ b/src/com/android/tradefed/result/proto/StreamProtoReceiver.java
@@ -189,6 +189,11 @@
         return true;
     }
 
+    /** If needed to ensure consistent reporting, complete the events of the module. */
+    public void completeModuleEvents() {
+        mParser.completeModuleEvents();
+    }
+
     private void parse(TestRecord receivedRecord) {
         try {
             TestLevel level = mParser.processNewProto(receivedRecord);
diff --git a/src/com/android/tradefed/sandbox/TradefedSandbox.java b/src/com/android/tradefed/sandbox/TradefedSandbox.java
index 14e4f6f..9718642 100644
--- a/src/com/android/tradefed/sandbox/TradefedSandbox.java
+++ b/src/com/android/tradefed/sandbox/TradefedSandbox.java
@@ -169,6 +169,9 @@
             result.setStderr(
                     String.format("Event receiver thread did not complete.:\n%s", stderrText));
         }
+        if (mProtoReceiver != null) {
+            mProtoReceiver.completeModuleEvents();
+        }
         PrettyPrintDelimiter.printStageDelimiter(
                 String.format(
                         "Execution of the tests occurred in the sandbox, you can find its logs "
diff --git a/tests/src/com/android/tradefed/result/proto/StreamProtoResultReporterTest.java b/tests/src/com/android/tradefed/result/proto/StreamProtoResultReporterTest.java
index 007df56..640940b 100644
--- a/tests/src/com/android/tradefed/result/proto/StreamProtoResultReporterTest.java
+++ b/tests/src/com/android/tradefed/result/proto/StreamProtoResultReporterTest.java
@@ -15,6 +15,7 @@
  */
 package com.android.tradefed.result.proto;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 
 import com.android.tradefed.config.ConfigurationDescriptor;
@@ -30,6 +31,7 @@
 import com.android.tradefed.testtype.suite.ModuleDefinition;
 import com.android.tradefed.util.proto.TfMetricProtoUtil;
 
+import org.easymock.Capture;
 import org.easymock.EasyMock;
 import org.junit.Before;
 import org.junit.Test;
@@ -188,6 +190,42 @@
         assertNull(receiver.getError());
     }
 
+    @Test
+    public void testStream_incompleteModule() throws Exception {
+        StreamProtoReceiver receiver =
+                new StreamProtoReceiver(mMockListener, mMainInvocationContext, true);
+        OptionSetter setter = new OptionSetter(mReporter);
+        Capture<FailureDescription> capture = new Capture<>();
+        try {
+            setter.setOptionValue(
+                    "proto-report-port", Integer.toString(receiver.getSocketServerPort()));
+            // Verify mocks
+            mMockListener.invocationStarted(EasyMock.anyObject());
+
+            mMockListener.testModuleStarted(EasyMock.anyObject());
+            mMockListener.testRunStarted(EasyMock.eq("arm64 module1"), EasyMock.eq(0));
+            mMockListener.testRunFailed(EasyMock.capture(capture));
+            mMockListener.testRunEnded(
+                    EasyMock.anyLong(), EasyMock.<HashMap<String, Metric>>anyObject());
+            mMockListener.testModuleEnded();
+
+            EasyMock.replay(mMockListener);
+            mReporter.invocationStarted(mInvocationContext);
+            // Run modules
+            mReporter.testModuleStarted(createModuleContext("arm64 module1"));
+            // It stops unexpectedly
+        } finally {
+            receiver.joinReceiver(2000);
+            receiver.close();
+            receiver.completeModuleEvents();
+        }
+        EasyMock.verify(mMockListener);
+        assertNull(receiver.getError());
+        assertEquals(
+                "Module was interrupted after starting, results are incomplete.",
+                capture.getValue().getErrorMessage());
+    }
+
     /** Helper to create a module context. */
     private IInvocationContext createModuleContext(String moduleId) {
         IInvocationContext context = new InvocationContext();