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