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.ChildReference;
22 import com.android.tradefed.result.proto.TestRecordProto.TestRecord;
23 import com.android.tradefed.util.FileUtil;
24 import com.android.tradefed.util.StreamUtil;
25 
26 import java.io.File;
27 import java.io.FileOutputStream;
28 import java.io.IOException;
29 
30 /** Proto reporter that dumps the {@link TestRecord} into a file. */
31 public class FileProtoResultReporter extends ProtoResultReporter {
32 
33     public static final String USE_DELIMITED_API = "use-delimited-api";
34 
35     @Option(
36             name = USE_DELIMITED_API,
37             description = "Use Proto.useDelimitedApi to save proto, otherwise use default api.")
38     private boolean mUseDelimitedApi = true;
39 
40     public static final String PROTO_OUTPUT_FILE = "proto-output-file";
41 
42     @Option(
43         name = PROTO_OUTPUT_FILE,
44         description = "File where the proto output will be saved. If unset, reporter will be inop."
45     )
46     private File mOutputFile = null;
47 
48     public static final String PERIODIC_PROTO_WRITING_OPTION = "periodic-proto-writing";
49 
50     @Option(
51         name = PERIODIC_PROTO_WRITING_OPTION,
52         description =
53                 "Whether or not to output intermediate proto per module following a numbered "
54                         + "sequence."
55     )
56     private boolean mPeriodicWriting = false;
57 
58     // Current index of the sequence of proto output
59     private int mIndex = 0;
60 
61     @Override
processStartInvocation( TestRecord invocationStartRecord, IInvocationContext invocationContext)62     public void processStartInvocation(
63             TestRecord invocationStartRecord, IInvocationContext invocationContext) {
64         writeProto(invocationStartRecord);
65     }
66 
67     @Override
processTestModuleEnd(TestRecord moduleRecord)68     public void processTestModuleEnd(TestRecord moduleRecord) {
69         writeProto(moduleRecord);
70     }
71 
72     @Override
processTestRunEnded(TestRecord runRecord, boolean moduleInProgress)73     public void processTestRunEnded(TestRecord runRecord, boolean moduleInProgress) {
74         if (!moduleInProgress) {
75             // If it's a testRun outside of the module scope, output it to ensure we support
76             // non-module use cases.
77             writeProto(runRecord);
78         }
79     }
80 
81     @Override
processFinalProto(TestRecord finalRecord)82     public void processFinalProto(TestRecord finalRecord) {
83         writeProto(finalRecord);
84     }
85 
86     @Override
createModuleChildReference(TestRecord record)87     protected ChildReference createModuleChildReference(TestRecord record) {
88         // Do not keep a copy of module record in invocation level
89         if (isPeriodicWriting()) {
90             return null;
91         }
92         return super.createModuleChildReference(record);
93     }
94 
95     /** Sets the file where to output the result. */
setFileOutput(File output)96     public void setFileOutput(File output) {
97         mOutputFile = output;
98     }
99 
100     /** Enable writing each module individualy to a file. */
setPeriodicWriting(boolean enabled)101     public void setPeriodicWriting(boolean enabled) {
102         mPeriodicWriting = enabled;
103     }
104 
105     /** Whether or not periodic writing is enabled. */
isPeriodicWriting()106     public boolean isPeriodicWriting() {
107         return mPeriodicWriting;
108     }
109 
writeProto(TestRecord record)110     private void writeProto(TestRecord record) {
111         if (mOutputFile == null) {
112             return;
113         }
114         FileOutputStream output = null;
115         File tmpFile = null;
116         try {
117             tmpFile = FileUtil.createTempFile("tmp-proto", "", mOutputFile.getParentFile());
118             File outputFile = mOutputFile;
119             if (mPeriodicWriting) {
120                 outputFile = new File(mOutputFile.getAbsolutePath() + mIndex);
121             }
122             // Write to the tmp file
123             output = new FileOutputStream(tmpFile);
124             if (mUseDelimitedApi) {
125                 record.writeDelimitedTo(output);
126             } else {
127                 record.writeTo(output);
128             }
129             if (mPeriodicWriting) {
130                 nextOutputFile();
131             }
132             // Move the tmp file to the new name when done writing.
133             tmpFile.renameTo(outputFile);
134         } catch (IOException e) {
135             CLog.e(e);
136             throw new RuntimeException(e);
137         } finally {
138             StreamUtil.close(output);
139         }
140     }
141 
nextOutputFile()142     private void nextOutputFile() {
143         mIndex++;
144     }
145 
setOutputFile(File outputFile)146     public void setOutputFile(File outputFile) {
147         mOutputFile = outputFile;
148     }
149 
getOutputFile()150     public File getOutputFile() {
151         return mOutputFile;
152     }
153 
setDelimitedOutput(boolean delimitedOutput)154     public void setDelimitedOutput(boolean delimitedOutput) {
155         mUseDelimitedApi = delimitedOutput;
156     }
157 }
158