1 /*
2  * Copyright (C) 2018 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package com.android.tradefed.result.proto;
17 
18 import com.android.tradefed.config.Option;
19 import com.android.tradefed.invoker.IInvocationContext;
20 import com.android.tradefed.log.LogUtil.CLog;
21 import com.android.tradefed.result.proto.TestRecordProto.TestRecord;
22 import com.android.tradefed.util.FileUtil;
23 import com.android.tradefed.util.StreamUtil;
24 
25 import java.io.File;
26 import java.io.FileOutputStream;
27 import java.io.IOException;
28 
29 /** Proto reporter that dumps the {@link TestRecord} into a file. */
30 public class FileProtoResultReporter extends ProtoResultReporter {
31 
32     public static final String USE_DELIMITED_API = "use-delimited-api";
33 
34     @Option(
35             name = USE_DELIMITED_API,
36             description = "Use Proto.useDelimitedApi to save proto, otherwise use default api.")
37     private boolean mUseDelimitedApi = true;
38 
39     public static final String PROTO_OUTPUT_FILE = "proto-output-file";
40 
41     @Option(
42         name = PROTO_OUTPUT_FILE,
43         description = "File where the proto output will be saved. If unset, reporter will be inop."
44     )
45     private File mOutputFile = null;
46 
47     public static final String PERIODIC_PROTO_WRITING_OPTION = "periodic-proto-writing";
48 
49     @Option(
50         name = PERIODIC_PROTO_WRITING_OPTION,
51         description =
52                 "Whether or not to output intermediate proto per module following a numbered "
53                         + "sequence."
54     )
55     private boolean mPeriodicWriting = false;
56 
57     // Current index of the sequence of proto output
58     private int mIndex = 0;
59 
60     @Override
processStartInvocation( TestRecord invocationStartRecord, IInvocationContext invocationContext)61     public void processStartInvocation(
62             TestRecord invocationStartRecord, IInvocationContext invocationContext) {
63         writeProto(invocationStartRecord);
64     }
65 
66     @Override
processTestModuleEnd(TestRecord moduleRecord)67     public void processTestModuleEnd(TestRecord moduleRecord) {
68         writeProto(moduleRecord);
69     }
70 
71     @Override
processTestRunEnded(TestRecord runRecord, boolean moduleInProgress)72     public void processTestRunEnded(TestRecord runRecord, boolean moduleInProgress) {
73         if (!moduleInProgress) {
74             // If it's a testRun outside of the module scope, output it to ensure we support
75             // non-module use cases.
76             writeProto(runRecord);
77         }
78     }
79 
80     @Override
processFinalProto(TestRecord finalRecord)81     public void processFinalProto(TestRecord finalRecord) {
82         writeProto(finalRecord);
83     }
84 
85     /** Sets the file where to output the result. */
setFileOutput(File output)86     public void setFileOutput(File output) {
87         mOutputFile = output;
88     }
89 
90     /** Enable writing each module individualy to a file. */
setPeriodicWriting(boolean enabled)91     public void setPeriodicWriting(boolean enabled) {
92         mPeriodicWriting = enabled;
93     }
94 
95     /** Whether or not periodic writing is enabled. */
isPeriodicWriting()96     public boolean isPeriodicWriting() {
97         return mPeriodicWriting;
98     }
99 
writeProto(TestRecord record)100     private void writeProto(TestRecord record) {
101         if (mOutputFile == null) {
102             return;
103         }
104         FileOutputStream output = null;
105         File tmpFile = null;
106         try {
107             tmpFile = FileUtil.createTempFile("tmp-proto", "", mOutputFile.getParentFile());
108             File outputFile = mOutputFile;
109             if (mPeriodicWriting) {
110                 outputFile = new File(mOutputFile.getAbsolutePath() + mIndex);
111             }
112             // Write to the tmp file
113             output = new FileOutputStream(tmpFile);
114             if (mUseDelimitedApi) {
115                 record.writeDelimitedTo(output);
116             } else {
117                 record.writeTo(output);
118             }
119             if (mPeriodicWriting) {
120                 nextOutputFile();
121             }
122             // Move the tmp file to the new name when done writing.
123             tmpFile.renameTo(outputFile);
124         } catch (IOException e) {
125             CLog.e(e);
126             throw new RuntimeException(e);
127         } finally {
128             StreamUtil.close(output);
129         }
130     }
131 
nextOutputFile()132     private void nextOutputFile() {
133         mIndex++;
134     }
135 }
136