1#!/usr/bin/env python3 2# Copyright (C) 2020 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 16import unittest 17 18from trace_processor.api import TraceProcessor, TraceProcessorException 19from trace_processor.protos import ProtoFactory 20 21 22class TestQueryResultIterator(unittest.TestCase): 23 # The numbers input into cells correspond the the CellType enum values 24 # defined under trace_processor.proto 25 CELL_VARINT = ProtoFactory().CellsBatch().CELL_VARINT 26 CELL_STRING = ProtoFactory().CellsBatch().CELL_STRING 27 CELL_INVALID = ProtoFactory().CellsBatch().CELL_INVALID 28 29 def test_one_batch(self): 30 int_values = [100, 200] 31 str_values = ['bar1', 'bar2'] 32 33 batch = ProtoFactory().CellsBatch() 34 batch.cells.extend([ 35 TestQueryResultIterator.CELL_STRING, 36 TestQueryResultIterator.CELL_VARINT, 37 TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT 38 ]) 39 batch.varint_cells.extend(int_values) 40 batch.string_cells = "\0".join(str_values) + "\0" 41 batch.is_last_batch = True 42 43 qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'], 44 [batch]) 45 46 for num, row in enumerate(qr_iterator): 47 self.assertEqual(row.foo_id, str_values[num]) 48 self.assertEqual(row.foo_num, int_values[num]) 49 50 def test_many_batches(self): 51 int_values = [100, 200, 300, 400] 52 str_values = ['bar1', 'bar2', 'bar3', 'bar4'] 53 54 batch_1 = ProtoFactory().CellsBatch() 55 batch_1.cells.extend([ 56 TestQueryResultIterator.CELL_STRING, 57 TestQueryResultIterator.CELL_VARINT, 58 TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT 59 ]) 60 batch_1.varint_cells.extend(int_values[:2]) 61 batch_1.string_cells = "\0".join(str_values[:2]) + "\0" 62 batch_1.is_last_batch = False 63 64 batch_2 = ProtoFactory().CellsBatch() 65 batch_2.cells.extend([ 66 TestQueryResultIterator.CELL_STRING, 67 TestQueryResultIterator.CELL_VARINT, 68 TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT 69 ]) 70 batch_2.varint_cells.extend(int_values[2:]) 71 batch_2.string_cells = "\0".join(str_values[2:]) + "\0" 72 batch_2.is_last_batch = True 73 74 qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'], 75 [batch_1, batch_2]) 76 77 for num, row in enumerate(qr_iterator): 78 self.assertEqual(row.foo_id, str_values[num]) 79 self.assertEqual(row.foo_num, int_values[num]) 80 81 def test_empty_batch(self): 82 batch = ProtoFactory().CellsBatch() 83 batch.is_last_batch = True 84 85 qr_iterator = TraceProcessor.QueryResultIterator([], [batch]) 86 87 for num, row in enumerate(qr_iterator): 88 self.assertIsNone(row.foo_id) 89 self.assertIsNone(row.foo_num) 90 91 def test_invalid_batch(self): 92 batch = ProtoFactory().CellsBatch() 93 94 qr_iterator = TraceProcessor.QueryResultIterator([], [batch]) 95 96 # Since the batch isn't defined as the last batch, the QueryResultsIterator 97 # expects another batch and thus raises IndexError as no next batch exists. 98 with self.assertRaises(IndexError): 99 for row in qr_iterator: 100 pass 101 102 def test_incorrect_cells_batch(self): 103 str_values = ['bar1', 'bar2'] 104 105 batch = ProtoFactory().CellsBatch() 106 batch.cells.extend([ 107 TestQueryResultIterator.CELL_STRING, 108 TestQueryResultIterator.CELL_VARINT, 109 TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT 110 ]) 111 batch.string_cells = "\0".join(str_values) + "\0" 112 batch.is_last_batch = True 113 114 qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'], 115 [batch]) 116 117 # The batch specifies there ought to be 2 cells of type VARINT and 2 cells 118 # of type STRING, but there are no string cells defined in the batch. Thus 119 # an IndexError occurs as it tries to access the empty string cells list. 120 with self.assertRaises(IndexError): 121 for row in qr_iterator: 122 pass 123 124 def test_incorrect_columns_batch(self): 125 batch = ProtoFactory().CellsBatch() 126 batch.cells.extend([ 127 TestQueryResultIterator.CELL_VARINT, TestQueryResultIterator.CELL_VARINT 128 ]) 129 batch.varint_cells.extend([100, 200]) 130 batch.is_last_batch = True 131 132 qr_iterator = TraceProcessor.QueryResultIterator( 133 ['foo_id', 'foo_num', 'foo_dur', 'foo_ms'], [batch]) 134 135 # It's always the case that the number of cells is a multiple of the number 136 # of columns. However, here this is clearly not the case, so when the 137 # iterator tries to access the cell for the third column, it raises an 138 # IndexError due to having exhausted the cells list. 139 with self.assertRaises(IndexError): 140 for row in qr_iterator: 141 pass 142 143 def test_invalid_cell_type(self): 144 batch = ProtoFactory().CellsBatch() 145 batch.cells.extend([ 146 TestQueryResultIterator.CELL_INVALID, 147 TestQueryResultIterator.CELL_VARINT 148 ]) 149 batch.varint_cells.extend([100, 200]) 150 batch.is_last_batch = True 151 152 qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'], 153 [batch]) 154 155 # In this batch we declare the columns types to be CELL_INVALID, 156 # CELL_VARINT but that doesn't match the data which are both ints* 157 # so we should raise a TraceProcessorException. 158 with self.assertRaises(TraceProcessorException): 159 for row in qr_iterator: 160 pass 161 162 def test_one_batch_as_pandas(self): 163 int_values = [100, 200] 164 str_values = ['bar1', 'bar2'] 165 166 batch = ProtoFactory().CellsBatch() 167 batch.cells.extend([ 168 TestQueryResultIterator.CELL_STRING, 169 TestQueryResultIterator.CELL_VARINT, 170 TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT 171 ]) 172 batch.varint_cells.extend(int_values) 173 batch.string_cells = "\0".join(str_values) + "\0" 174 batch.is_last_batch = True 175 176 qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'], 177 [batch]) 178 179 qr_df = qr_iterator.as_pandas_dataframe() 180 for num, row in qr_df.iterrows(): 181 self.assertEqual(row['foo_id'], str_values[num]) 182 self.assertEqual(row['foo_num'], int_values[num]) 183 184 def test_many_batches_as_pandas(self): 185 int_values = [100, 200, 300, 400] 186 str_values = ['bar1', 'bar2', 'bar3', 'bar4'] 187 188 batch_1 = ProtoFactory().CellsBatch() 189 batch_1.cells.extend([ 190 TestQueryResultIterator.CELL_STRING, 191 TestQueryResultIterator.CELL_VARINT, 192 TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT 193 ]) 194 batch_1.varint_cells.extend(int_values[:2]) 195 batch_1.string_cells = "\0".join(str_values[:2]) + "\0" 196 batch_1.is_last_batch = False 197 198 batch_2 = ProtoFactory().CellsBatch() 199 batch_2.cells.extend([ 200 TestQueryResultIterator.CELL_STRING, 201 TestQueryResultIterator.CELL_VARINT, 202 TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT 203 ]) 204 batch_2.varint_cells.extend(int_values[2:]) 205 batch_2.string_cells = "\0".join(str_values[2:]) + "\0" 206 batch_2.is_last_batch = True 207 208 qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'], 209 [batch_1, batch_2]) 210 211 qr_df = qr_iterator.as_pandas_dataframe() 212 for num, row in qr_df.iterrows(): 213 self.assertEqual(row['foo_id'], str_values[num]) 214 self.assertEqual(row['foo_num'], int_values[num]) 215 216 def test_empty_batch_as_pandas(self): 217 batch = ProtoFactory().CellsBatch() 218 batch.is_last_batch = True 219 220 qr_iterator = TraceProcessor.QueryResultIterator([], [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 227 def test_invalid_batch_as_pandas(self): 228 batch = ProtoFactory().CellsBatch() 229 230 qr_iterator = TraceProcessor.QueryResultIterator([], [batch]) 231 232 # Since the batch isn't defined as the last batch, the QueryResultsIterator 233 # expects another batch and thus raises IndexError as no next batch exists. 234 with self.assertRaises(IndexError): 235 qr_df = qr_iterator.as_pandas_dataframe() 236 237 def test_incorrect_cells_batch_as_pandas(self): 238 str_values = ['bar1', 'bar2'] 239 240 batch = ProtoFactory().CellsBatch() 241 batch.cells.extend([ 242 TestQueryResultIterator.CELL_STRING, 243 TestQueryResultIterator.CELL_VARINT, 244 TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT 245 ]) 246 batch.string_cells = "\0".join(str_values) + "\0" 247 batch.is_last_batch = True 248 249 qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'], 250 [batch]) 251 252 # The batch specifies there ought to be 2 cells of type VARINT and 2 cells 253 # of type STRING, but there are no string cells defined in the batch. Thus 254 # an IndexError occurs as it tries to access the empty string cells list. 255 with self.assertRaises(IndexError): 256 qr_df = qr_iterator.as_pandas_dataframe() 257 258 def test_incorrect_columns_batch_as_pandas(self): 259 batch = ProtoFactory().CellsBatch() 260 batch.cells.extend([ 261 TestQueryResultIterator.CELL_VARINT, TestQueryResultIterator.CELL_VARINT 262 ]) 263 batch.varint_cells.extend([100, 200]) 264 batch.is_last_batch = True 265 266 qr_iterator = TraceProcessor.QueryResultIterator( 267 ['foo_id', 'foo_num', 'foo_dur', 'foo_ms'], [batch]) 268 269 # It's always the case that the number of cells is a multiple of the number 270 # of columns. However, here this is clearly not the case, so when the 271 # iterator tries to access the cell for the third column, it raises an 272 # IndexError due to having exhausted the cells list. 273 with self.assertRaises(IndexError): 274 qr_df = qr_iterator.as_pandas_dataframe() 275 276 def test_invalid_cell_type_as_pandas(self): 277 batch = ProtoFactory().CellsBatch() 278 batch.cells.extend([ 279 TestQueryResultIterator.CELL_INVALID, 280 TestQueryResultIterator.CELL_VARINT 281 ]) 282 batch.varint_cells.extend([100, 200]) 283 batch.is_last_batch = True 284 285 qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'], 286 [batch]) 287 288 # In this batch we declare the columns types to be CELL_INVALID, 289 # CELL_VARINT but that doesn't match the data which are both ints* 290 # so we should raise a TraceProcessorException. 291 with self.assertRaises(TraceProcessorException): 292 qr_df = qr_iterator.as_pandas_dataframe() 293