1# Copyright (C) 2020 The Android Open Source Project 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15import unittest 16 17from perfetto.trace_processor.api import TraceProcessor 18from perfetto.trace_processor.api import TraceProcessorException 19from perfetto.trace_processor.api import PLATFORM_DELEGATE 20from perfetto.trace_processor.protos import ProtoFactory 21 22PROTO_FACTORY = ProtoFactory(PLATFORM_DELEGATE()) 23 24 25class TestQueryResultIterator(unittest.TestCase): 26 # The numbers input into cells correspond the CellType enum values 27 # defined under trace_processor.proto 28 CELL_VARINT = PROTO_FACTORY.CellsBatch().CELL_VARINT 29 CELL_STRING = PROTO_FACTORY.CellsBatch().CELL_STRING 30 CELL_INVALID = PROTO_FACTORY.CellsBatch().CELL_INVALID 31 CELL_NULL = PROTO_FACTORY.CellsBatch().CELL_NULL 32 33 def test_one_batch(self): 34 int_values = [100, 200] 35 str_values = ['bar1', 'bar2'] 36 37 batch = PROTO_FACTORY.CellsBatch() 38 batch.cells.extend([ 39 TestQueryResultIterator.CELL_STRING, 40 TestQueryResultIterator.CELL_VARINT, 41 TestQueryResultIterator.CELL_NULL, 42 TestQueryResultIterator.CELL_STRING, 43 TestQueryResultIterator.CELL_VARINT, 44 TestQueryResultIterator.CELL_NULL, 45 ]) 46 batch.varint_cells.extend(int_values) 47 batch.string_cells = "\0".join(str_values) + "\0" 48 batch.is_last_batch = True 49 50 qr_iterator = TraceProcessor.QueryResultIterator( 51 ['foo_id', 'foo_num', 'foo_null'], [batch]) 52 53 for num, row in enumerate(qr_iterator): 54 self.assertEqual(row.foo_id, str_values[num]) 55 self.assertEqual(row.foo_num, int_values[num]) 56 self.assertEqual(row.foo_null, None) 57 58 def test_many_batches(self): 59 int_values = [100, 200, 300, 400] 60 str_values = ['bar1', 'bar2', 'bar3', 'bar4'] 61 62 batch_1 = PROTO_FACTORY.CellsBatch() 63 batch_1.cells.extend([ 64 TestQueryResultIterator.CELL_STRING, 65 TestQueryResultIterator.CELL_VARINT, 66 TestQueryResultIterator.CELL_NULL, 67 TestQueryResultIterator.CELL_STRING, 68 TestQueryResultIterator.CELL_VARINT, 69 TestQueryResultIterator.CELL_NULL, 70 ]) 71 batch_1.varint_cells.extend(int_values[:2]) 72 batch_1.string_cells = "\0".join(str_values[:2]) + "\0" 73 batch_1.is_last_batch = False 74 75 batch_2 = PROTO_FACTORY.CellsBatch() 76 batch_2.cells.extend([ 77 TestQueryResultIterator.CELL_STRING, 78 TestQueryResultIterator.CELL_VARINT, 79 TestQueryResultIterator.CELL_NULL, 80 TestQueryResultIterator.CELL_STRING, 81 TestQueryResultIterator.CELL_VARINT, 82 TestQueryResultIterator.CELL_NULL, 83 ]) 84 batch_2.varint_cells.extend(int_values[2:]) 85 batch_2.string_cells = "\0".join(str_values[2:]) + "\0" 86 batch_2.is_last_batch = True 87 88 qr_iterator = TraceProcessor.QueryResultIterator( 89 ['foo_id', 'foo_num', 'foo_null'], [batch_1, batch_2]) 90 91 for num, row in enumerate(qr_iterator): 92 self.assertEqual(row.foo_id, str_values[num]) 93 self.assertEqual(row.foo_num, int_values[num]) 94 self.assertEqual(row.foo_null, None) 95 96 def test_empty_batch(self): 97 batch = PROTO_FACTORY.CellsBatch() 98 batch.is_last_batch = True 99 100 qr_iterator = TraceProcessor.QueryResultIterator([], [batch]) 101 102 for num, row in enumerate(qr_iterator): 103 self.assertIsNone(row.foo_id) 104 self.assertIsNone(row.foo_num) 105 106 def test_invalid_batch(self): 107 batch = PROTO_FACTORY.CellsBatch() 108 109 # Since the batch isn't defined as the last batch, the QueryResultsIterator 110 # expects another batch and thus raises IndexError as no next batch exists. 111 with self.assertRaises(IndexError): 112 qr_iterator = TraceProcessor.QueryResultIterator([], [batch]) 113 114 def test_null_cells(self): 115 int_values = [100, 200, 300, 500, 600] 116 str_values = ['bar1', 'bar2', 'bar3'] 117 118 batch = PROTO_FACTORY.CellsBatch() 119 batch.cells.extend([ 120 TestQueryResultIterator.CELL_STRING, 121 TestQueryResultIterator.CELL_VARINT, 122 TestQueryResultIterator.CELL_VARINT, 123 TestQueryResultIterator.CELL_STRING, 124 TestQueryResultIterator.CELL_VARINT, 125 TestQueryResultIterator.CELL_NULL, 126 TestQueryResultIterator.CELL_STRING, 127 TestQueryResultIterator.CELL_VARINT, 128 TestQueryResultIterator.CELL_VARINT, 129 ]) 130 batch.varint_cells.extend(int_values) 131 batch.string_cells = "\0".join(str_values) + "\0" 132 batch.is_last_batch = True 133 134 qr_iterator = TraceProcessor.QueryResultIterator( 135 ['foo_id', 'foo_num', 'foo_num_2'], [batch]) 136 137 # Any cell (and thus column in a row) can be set to null 138 # In this query result, foo_num_2 of row 2 was set to null 139 # Test to see that all the rows are still returned correctly 140 int_values_check = [100, 200, 300, None, 500, 600] 141 for num, row in enumerate(qr_iterator): 142 self.assertEqual(row.foo_id, str_values[num]) 143 self.assertEqual(row.foo_num, int_values_check[num * 2]) 144 self.assertEqual(row.foo_num_2, int_values_check[num * 2 + 1]) 145 146 def test_incorrect_cells_batch(self): 147 str_values = ['bar1', 'bar2'] 148 149 batch = PROTO_FACTORY.CellsBatch() 150 batch.cells.extend([ 151 TestQueryResultIterator.CELL_STRING, 152 TestQueryResultIterator.CELL_VARINT, 153 TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT 154 ]) 155 batch.string_cells = "\0".join(str_values) + "\0" 156 batch.is_last_batch = True 157 158 qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'], 159 [batch]) 160 161 # The batch specifies there ought to be 2 cells of type VARINT and 2 cells 162 # of type STRING, but there are no string cells defined in the batch. Thus 163 # an IndexError occurs as it tries to access the empty string cells list. 164 with self.assertRaises(IndexError): 165 for row in qr_iterator: 166 pass 167 168 def test_incorrect_columns_batch(self): 169 batch = PROTO_FACTORY.CellsBatch() 170 batch.cells.extend([ 171 TestQueryResultIterator.CELL_VARINT, TestQueryResultIterator.CELL_VARINT 172 ]) 173 batch.varint_cells.extend([100, 200]) 174 batch.is_last_batch = True 175 176 # It's always the case that the number of cells is a multiple of the number 177 # of columns. However, here this is clearly not the case, so raise a 178 # TraceProcessorException during the data integrity check in the constructor 179 with self.assertRaises(TraceProcessorException): 180 qr_iterator = TraceProcessor.QueryResultIterator( 181 ['foo_id', 'foo_num', 'foo_dur', 'foo_ms'], [batch]) 182 183 def test_invalid_cell_type(self): 184 batch = PROTO_FACTORY.CellsBatch() 185 batch.cells.extend([ 186 TestQueryResultIterator.CELL_INVALID, 187 TestQueryResultIterator.CELL_VARINT 188 ]) 189 batch.varint_cells.extend([100, 200]) 190 batch.is_last_batch = True 191 192 qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'], 193 [batch]) 194 195 # In this batch we declare the columns types to be CELL_INVALID, 196 # CELL_VARINT but that doesn't match the data which are both ints* 197 # so we should raise a TraceProcessorException. 198 with self.assertRaises(TraceProcessorException): 199 for row in qr_iterator: 200 pass 201 202 def test_one_batch_as_pandas(self): 203 int_values = [100, 200] 204 str_values = ['bar1', 'bar2'] 205 206 batch = PROTO_FACTORY.CellsBatch() 207 batch.cells.extend([ 208 TestQueryResultIterator.CELL_STRING, 209 TestQueryResultIterator.CELL_VARINT, 210 TestQueryResultIterator.CELL_NULL, 211 TestQueryResultIterator.CELL_STRING, 212 TestQueryResultIterator.CELL_VARINT, 213 TestQueryResultIterator.CELL_NULL, 214 ]) 215 batch.varint_cells.extend(int_values) 216 batch.string_cells = "\0".join(str_values) + "\0" 217 batch.is_last_batch = True 218 219 qr_iterator = TraceProcessor.QueryResultIterator( 220 ['foo_id', 'foo_num', 'foo_null'], [batch]) 221 222 qr_df = qr_iterator.as_pandas_dataframe() 223 for num, row in qr_df.iterrows(): 224 self.assertEqual(row['foo_id'], str_values[num]) 225 self.assertEqual(row['foo_num'], int_values[num]) 226 self.assertEqual(row['foo_null'], None) 227 228 def test_many_batches_as_pandas(self): 229 int_values = [100, 200, 300, 400] 230 str_values = ['bar1', 'bar2', 'bar3', 'bar4'] 231 232 batch_1 = PROTO_FACTORY.CellsBatch() 233 batch_1.cells.extend([ 234 TestQueryResultIterator.CELL_STRING, 235 TestQueryResultIterator.CELL_VARINT, 236 TestQueryResultIterator.CELL_NULL, 237 TestQueryResultIterator.CELL_STRING, 238 TestQueryResultIterator.CELL_VARINT, 239 TestQueryResultIterator.CELL_NULL, 240 ]) 241 batch_1.varint_cells.extend(int_values[:2]) 242 batch_1.string_cells = "\0".join(str_values[:2]) + "\0" 243 batch_1.is_last_batch = False 244 245 batch_2 = PROTO_FACTORY.CellsBatch() 246 batch_2.cells.extend([ 247 TestQueryResultIterator.CELL_STRING, 248 TestQueryResultIterator.CELL_VARINT, 249 TestQueryResultIterator.CELL_NULL, 250 TestQueryResultIterator.CELL_STRING, 251 TestQueryResultIterator.CELL_VARINT, 252 TestQueryResultIterator.CELL_NULL, 253 ]) 254 batch_2.varint_cells.extend(int_values[2:]) 255 batch_2.string_cells = "\0".join(str_values[2:]) + "\0" 256 batch_2.is_last_batch = True 257 258 qr_iterator = TraceProcessor.QueryResultIterator( 259 ['foo_id', 'foo_num', 'foo_null'], [batch_1, batch_2]) 260 261 qr_df = qr_iterator.as_pandas_dataframe() 262 for num, row in qr_df.iterrows(): 263 self.assertEqual(row['foo_id'], str_values[num]) 264 self.assertEqual(row['foo_num'], int_values[num]) 265 self.assertEqual(row['foo_null'], None) 266 267 def test_empty_batch_as_pandas(self): 268 batch = PROTO_FACTORY.CellsBatch() 269 batch.is_last_batch = True 270 271 qr_iterator = TraceProcessor.QueryResultIterator([], [batch]) 272 273 qr_df = qr_iterator.as_pandas_dataframe() 274 for num, row in qr_df.iterrows(): 275 self.assertEqual(row['foo_id'], str_values[num]) 276 self.assertEqual(row['foo_num'], int_values[num]) 277 278 def test_null_cells_as_pandas(self): 279 int_values = [100, 200, 300, 500, 600] 280 str_values = ['bar1', 'bar2', 'bar3'] 281 282 batch = PROTO_FACTORY.CellsBatch() 283 batch.cells.extend([ 284 TestQueryResultIterator.CELL_STRING, 285 TestQueryResultIterator.CELL_VARINT, 286 TestQueryResultIterator.CELL_VARINT, 287 TestQueryResultIterator.CELL_STRING, 288 TestQueryResultIterator.CELL_VARINT, 289 TestQueryResultIterator.CELL_NULL, 290 TestQueryResultIterator.CELL_STRING, 291 TestQueryResultIterator.CELL_VARINT, 292 TestQueryResultIterator.CELL_VARINT, 293 ]) 294 batch.varint_cells.extend(int_values) 295 batch.string_cells = "\0".join(str_values) + "\0" 296 batch.is_last_batch = True 297 298 qr_iterator = TraceProcessor.QueryResultIterator( 299 ['foo_id', 'foo_num', 'foo_num_2'], [batch]) 300 qr_df = qr_iterator.as_pandas_dataframe() 301 302 # Any cell (and thus column in a row) can be set to null 303 # In this query result, foo_num_2 of row 2 was set to null 304 # Test to see that all the rows are still returned correctly 305 int_values_check = [100, 200, 300, None, 500, 600] 306 for num, row in qr_df.iterrows(): 307 self.assertEqual(row['foo_id'], str_values[num]) 308 self.assertEqual(row['foo_num'], int_values_check[num * 2]) 309 self.assertEqual(row['foo_num_2'], int_values_check[num * 2 + 1]) 310 311 def test_incorrect_cells_batch_as_pandas(self): 312 str_values = ['bar1', 'bar2'] 313 314 batch = PROTO_FACTORY.CellsBatch() 315 batch.cells.extend([ 316 TestQueryResultIterator.CELL_STRING, 317 TestQueryResultIterator.CELL_VARINT, 318 TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT 319 ]) 320 batch.string_cells = "\0".join(str_values) + "\0" 321 batch.is_last_batch = True 322 323 qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'], 324 [batch]) 325 326 # The batch specifies there ought to be 2 cells of type VARINT and 2 cells 327 # of type STRING, but there are no string cells defined in the batch. Thus 328 # an IndexError occurs as it tries to access the empty string cells list. 329 with self.assertRaises(IndexError): 330 _ = qr_iterator.as_pandas_dataframe() 331 332 def test_invalid_cell_type_as_pandas(self): 333 batch = PROTO_FACTORY.CellsBatch() 334 batch.cells.extend([ 335 TestQueryResultIterator.CELL_INVALID, 336 TestQueryResultIterator.CELL_VARINT 337 ]) 338 batch.varint_cells.extend([100, 200]) 339 batch.is_last_batch = True 340 341 qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'], 342 [batch]) 343 344 # In this batch we declare the columns types to be CELL_INVALID, 345 # CELL_VARINT but that doesn't match the data which are both ints* 346 # so we should raise a TraceProcessorException. 347 with self.assertRaises(TraceProcessorException): 348 _ = qr_iterator.as_pandas_dataframe() 349