| # Copyright 2011 The ChromiumOS Authors |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """Unit tests for the table module.""" |
| |
| import tempfile |
| |
| from chromite.lib import cros_test_lib |
| from chromite.lib import table |
| from chromite.utils import outcap |
| |
| |
| # pylint: disable=protected-access |
| |
| |
| class TableTest(cros_test_lib.TempDirTestCase): |
| """Unit tests for the Table class.""" |
| |
| COL0 = "Column1" |
| COL1 = "Column2" |
| COL2 = "Column3" |
| COL3 = "Column4" |
| COLUMNS = [COL0, COL1, COL2, COL3] |
| |
| ROW0 = {COL0: "Xyz", COL1: "Bcd", COL2: "Cde"} |
| ROW1 = {COL0: "Abc", COL1: "Bcd", COL2: "Opq", COL3: "Foo"} |
| ROW2 = {COL0: "Abc", COL1: "Nop", COL2: "Wxy", COL3: "Bar"} |
| |
| EXTRAROW = {COL1: "Walk", COL2: "The", COL3: "Line"} |
| EXTRAROWOUT = {COL0: "", COL1: "Walk", COL2: "The", COL3: "Line"} |
| |
| ROW0a = {COL0: "Xyz", COL1: "Bcd", COL2: "Cde", COL3: "Yay"} |
| ROW0b = {COL0: "Xyz", COL1: "Bcd", COL2: "Cde", COL3: "Boo"} |
| ROW1a = {COL0: "Abc", COL1: "Bcd", COL2: "Opq", COL3: "Blu"} |
| |
| EXTRACOL = "ExtraCol" |
| EXTRACOLUMNS = [COL0, EXTRACOL, COL1, COL2] |
| |
| EROW0 = {COL0: "Xyz", EXTRACOL: "Yay", COL1: "Bcd", COL2: "Cde"} |
| EROW1 = {COL0: "Abc", EXTRACOL: "Hip", COL1: "Bcd", COL2: "Opq"} |
| EROW2 = {COL0: "Abc", EXTRACOL: "Yay", COL1: "Nop", COL2: "Wxy"} |
| |
| def _GetRowValsInOrder(self, row): |
| """Take |row| dict and return correctly ordered values in a list.""" |
| vals = [] |
| for col in self.COLUMNS: |
| vals.append(row.get(col, "")) |
| |
| return vals |
| |
| def _GetFullRowFor(self, row, cols): |
| return dict((col, row.get(col, "")) for col in cols) |
| |
| def assertRowsEqual(self, row1, row2): |
| # Determine column superset |
| cols = set(row1) | set(row2) |
| self.assertEqual( |
| self._GetFullRowFor(row1, cols), self._GetFullRowFor(row2, cols) |
| ) |
| |
| def assertRowListsEqual(self, rows1, rows2): |
| for row1, row2 in zip(rows1, rows2): |
| self.assertRowsEqual(row1, row2) |
| |
| def setUp(self): |
| self._table = self._CreateTableWithRows( |
| self.COLUMNS, [self.ROW0, self.ROW1, self.ROW2] |
| ) |
| |
| def _CreateTableWithRows(self, cols, rows): |
| mytable = table.Table(list(cols)) |
| if rows: |
| for row in rows: |
| mytable.AppendRow(dict(row)) |
| return mytable |
| |
| def testLen(self): |
| self.assertEqual(3, len(self._table)) |
| |
| def testGetNumRows(self): |
| self.assertEqual(3, self._table.GetNumRows()) |
| |
| def testGetNumColumns(self): |
| self.assertEqual(4, self._table.GetNumColumns()) |
| |
| def testGetColumns(self): |
| self.assertEqual(self.COLUMNS, self._table.GetColumns()) |
| |
| def testGetColumnIndex(self): |
| self.assertEqual(0, self._table.GetColumnIndex(self.COL0)) |
| self.assertEqual(1, self._table.GetColumnIndex(self.COL1)) |
| self.assertEqual(2, self._table.GetColumnIndex(self.COL2)) |
| |
| def testGetColumnByIndex(self): |
| self.assertEqual(self.COL0, self._table.GetColumnByIndex(0)) |
| self.assertEqual(self.COL1, self._table.GetColumnByIndex(1)) |
| self.assertEqual(self.COL2, self._table.GetColumnByIndex(2)) |
| |
| def testGetByIndex(self): |
| self.assertRowsEqual(self.ROW0, self._table.GetRowByIndex(0)) |
| self.assertRowsEqual(self.ROW0, self._table[0]) |
| |
| self.assertRowsEqual(self.ROW2, self._table.GetRowByIndex(2)) |
| self.assertRowsEqual(self.ROW2, self._table[2]) |
| |
| def testSlice(self): |
| self.assertRowListsEqual([self.ROW0, self.ROW1], self._table[0:2]) |
| self.assertRowListsEqual([self.ROW2], self._table[-1:]) |
| |
| def testGetByValue(self): |
| rows = self._table.GetRowsByValue({self.COL0: "Abc"}) |
| self.assertEqual([self.ROW1, self.ROW2], rows) |
| rows = self._table.GetRowsByValue({self.COL2: "Opq"}) |
| self.assertEqual([self.ROW1], rows) |
| rows = self._table.GetRowsByValue({self.COL3: "Foo"}) |
| self.assertEqual([self.ROW1], rows) |
| |
| def testGetIndicesByValue(self): |
| indices = self._table.GetRowIndicesByValue({self.COL0: "Abc"}) |
| self.assertEqual([1, 2], indices) |
| indices = self._table.GetRowIndicesByValue({self.COL2: "Opq"}) |
| self.assertEqual([1], indices) |
| indices = self._table.GetRowIndicesByValue({self.COL3: "Foo"}) |
| self.assertEqual([1], indices) |
| |
| def testAppendRowDict(self): |
| self._table.AppendRow(dict(self.EXTRAROW)) |
| self.assertEqual(4, self._table.GetNumRows()) |
| self.assertEqual(self.EXTRAROWOUT, self._table[len(self._table) - 1]) |
| |
| def testAppendRowList(self): |
| self._table.AppendRow(self._GetRowValsInOrder(self.EXTRAROW)) |
| self.assertEqual(4, self._table.GetNumRows()) |
| self.assertEqual(self.EXTRAROWOUT, self._table[len(self._table) - 1]) |
| |
| def testSetRowDictByIndex(self): |
| self._table.SetRowByIndex(1, dict(self.EXTRAROW)) |
| self.assertEqual(3, self._table.GetNumRows()) |
| self.assertEqual(self.EXTRAROWOUT, self._table[1]) |
| |
| def testSetRowListByIndex(self): |
| self._table.SetRowByIndex(1, self._GetRowValsInOrder(self.EXTRAROW)) |
| self.assertEqual(3, self._table.GetNumRows()) |
| self.assertEqual(self.EXTRAROWOUT, self._table[1]) |
| |
| def testRemoveRowByIndex(self): |
| self._table.RemoveRowByIndex(1) |
| self.assertEqual(2, self._table.GetNumRows()) |
| self.assertEqual(self.ROW2, self._table[1]) |
| |
| def testRemoveRowBySlice(self): |
| del self._table[0:2] |
| self.assertEqual(1, self._table.GetNumRows()) |
| self.assertEqual(self.ROW2, self._table[0]) |
| |
| def testIteration(self): |
| ix = 0 |
| for row in self._table: |
| self.assertEqual(row, self._table[ix]) |
| ix += 1 |
| |
| def testClear(self): |
| self._table.Clear() |
| self.assertEqual(0, len(self._table)) |
| |
| def testMergeRows(self): |
| # This merge should fail without a merge rule. Capture stderr to avoid |
| # scary error message in test output. |
| with outcap.OutputCapturer(): |
| self.assertRaises( |
| ValueError, self._table._MergeRow, self.ROW0a, self.COL0 |
| ) |
| |
| # Merge but stick with current row where different. |
| self._table._MergeRow( |
| self.ROW0a, self.COL0, merge_rules={self.COL3: "accept_this_val"} |
| ) |
| self.assertEqual(3, len(self._table)) |
| self.assertRowsEqual(self.ROW0, self._table[0]) |
| |
| # Merge and use new row where different. |
| self._table._MergeRow( |
| self.ROW0a, self.COL0, merge_rules={self.COL3: "accept_other_val"} |
| ) |
| self.assertEqual(3, len(self._table)) |
| self.assertRowsEqual(self.ROW0a, self._table[0]) |
| |
| # Merge and combine column values where different |
| self._table._MergeRow( |
| self.ROW1a, self.COL2, merge_rules={self.COL3: "join_with: "} |
| ) |
| self.assertEqual(3, len(self._table)) |
| final_row = dict(self.ROW1a) |
| final_row[self.COL3] = ( |
| self.ROW1[self.COL3] + " " + self.ROW1a[self.COL3] |
| ) |
| self.assertRowsEqual(final_row, self._table[1]) |
| |
| def testMergeTablesSameCols(self): |
| other_table = self._CreateTableWithRows( |
| self.COLUMNS, [self.ROW0b, self.ROW1a, self.ROW2] |
| ) |
| |
| self._table.MergeTable( |
| other_table, self.COL2, merge_rules={self.COL3: "join_with: "} |
| ) |
| |
| final_row0 = self.ROW0b |
| final_row1 = dict(self.ROW1a) |
| final_row1[self.COL3] = ( |
| self.ROW1[self.COL3] + " " + self.ROW1a[self.COL3] |
| ) |
| final_row2 = self.ROW2 |
| self.assertRowsEqual(final_row0, self._table[0]) |
| self.assertRowsEqual(final_row1, self._table[1]) |
| self.assertRowsEqual(final_row2, self._table[2]) |
| |
| def testMergeTablesNewCols(self): |
| self.assertFalse(self._table.HasColumn(self.EXTRACOL)) |
| |
| other_rows = [self.EROW0, self.EROW1, self.EROW2] |
| other_table = self._CreateTableWithRows(self.EXTRACOLUMNS, other_rows) |
| |
| self._table.MergeTable( |
| other_table, |
| self.COL2, |
| allow_new_columns=True, |
| merge_rules={self.COL3: "join_by_space"}, |
| ) |
| |
| self.assertTrue(self._table.HasColumn(self.EXTRACOL)) |
| self.assertEqual(5, self._table.GetNumColumns()) |
| self.assertEqual(1, self._table.GetColumnIndex(self.EXTRACOL)) |
| |
| final_row0 = dict(self.ROW0) |
| final_row0[self.EXTRACOL] = self.EROW0[self.EXTRACOL] |
| final_row1 = dict(self.ROW1) |
| final_row1[self.EXTRACOL] = self.EROW1[self.EXTRACOL] |
| final_row2 = dict(self.ROW2) |
| final_row2[self.EXTRACOL] = self.EROW2[self.EXTRACOL] |
| self.assertRowsEqual(final_row0, self._table[0]) |
| self.assertRowsEqual(final_row1, self._table[1]) |
| self.assertRowsEqual(final_row2, self._table[2]) |
| |
| def testSort1(self): |
| self.assertRowsEqual(self.ROW0, self._table[0]) |
| self.assertRowsEqual(self.ROW1, self._table[1]) |
| self.assertRowsEqual(self.ROW2, self._table[2]) |
| |
| # Sort by COL3 |
| self._table.Sort(lambda row: row[self.COL3]) |
| |
| self.assertEqual(3, len(self._table)) |
| self.assertRowsEqual(self.ROW0, self._table[0]) |
| self.assertRowsEqual(self.ROW2, self._table[1]) |
| self.assertRowsEqual(self.ROW1, self._table[2]) |
| |
| # Reverse sort by COL3 |
| self._table.Sort(lambda row: row[self.COL3], reverse=True) |
| |
| self.assertEqual(3, len(self._table)) |
| self.assertRowsEqual(self.ROW1, self._table[0]) |
| self.assertRowsEqual(self.ROW2, self._table[1]) |
| self.assertRowsEqual(self.ROW0, self._table[2]) |
| |
| def testSort2(self): |
| """Test multiple key sort.""" |
| self.assertRowsEqual(self.ROW0, self._table[0]) |
| self.assertRowsEqual(self.ROW1, self._table[1]) |
| self.assertRowsEqual(self.ROW2, self._table[2]) |
| |
| # Sort by COL0 then COL1 |
| def sorter(row): |
| return (row[self.COL0], row[self.COL1]) |
| |
| self._table.Sort(sorter) |
| |
| self.assertEqual(3, len(self._table)) |
| self.assertRowsEqual(self.ROW1, self._table[0]) |
| self.assertRowsEqual(self.ROW2, self._table[1]) |
| self.assertRowsEqual(self.ROW0, self._table[2]) |
| |
| # Reverse the sort |
| self._table.Sort(sorter, reverse=True) |
| |
| self.assertEqual(3, len(self._table)) |
| self.assertRowsEqual(self.ROW0, self._table[0]) |
| self.assertRowsEqual(self.ROW2, self._table[1]) |
| self.assertRowsEqual(self.ROW1, self._table[2]) |
| |
| def testSplitCSVLine(self): |
| """Test splitting of csv line.""" |
| tests = { |
| "a,b,c,d": ["a", "b", "c", "d"], |
| "a, b, c, d": ["a", " b", " c", " d"], |
| "a,b,c,": ["a", "b", "c", ""], |
| 'a,"b c",d': ["a", "b c", "d"], |
| 'a,"b, c",d': ["a", "b, c", "d"], |
| 'a,"b, c, d",e': ["a", "b, c, d", "e"], |
| 'a,"""b, c""",d': ["a", '"b, c"', "d"], |
| 'a,"""b, c"", d",e': ["a", '"b, c", d', "e"], |
| # Following not real Google Spreadsheet cases. |
| r"a,b\,c,d": ["a", "b,c", "d"], |
| 'a,",c': ["a", '",c'], |
| 'a,"",c': ["a", "", "c"], |
| } |
| # TODO(b/236161656): Fix. |
| # pylint: disable-next=consider-using-dict-items |
| for line in tests: |
| vals = table.Table._SplitCSVLine(line) |
| self.assertEqual(vals, tests[line]) |
| |
| def testWriteReadCSV(self): |
| """Write and Read CSV and verify contents preserved.""" |
| # This also tests the Table == and != operators. |
| _, path = tempfile.mkstemp(text=True) |
| with open(path, "w", encoding="utf-8") as tmpfile: |
| self._table.WriteCSV(tmpfile) |
| mytable = table.Table.LoadFromCSV(path) |
| self.assertEqual(mytable, self._table) |
| self.assertFalse(mytable != self._table) |
| |
| def testInsertColumn(self): |
| self._table.InsertColumn(1, self.EXTRACOL, "blah") |
| goldenrow = dict(self.ROW1) |
| goldenrow[self.EXTRACOL] = "blah" |
| self.assertRowsEqual(goldenrow, self._table.GetRowByIndex(1)) |
| self.assertEqual(self.EXTRACOL, self._table.GetColumnByIndex(1)) |
| |
| def testAppendColumn(self): |
| self._table.AppendColumn(self.EXTRACOL, "blah") |
| goldenrow = dict(self.ROW1) |
| goldenrow[self.EXTRACOL] = "blah" |
| self.assertRowsEqual(goldenrow, self._table.GetRowByIndex(1)) |
| col_size = self._table.GetNumColumns() |
| self.assertEqual( |
| self.EXTRACOL, self._table.GetColumnByIndex(col_size - 1) |
| ) |
| |
| def testProcessRows(self): |
| def Processor(row): |
| row[self.COL0] = row[self.COL0] + " processed" |
| |
| self._table.ProcessRows(Processor) |
| |
| final_row0 = dict(self.ROW0) |
| final_row0[self.COL0] += " processed" |
| final_row1 = dict(self.ROW1) |
| final_row1[self.COL0] += " processed" |
| final_row2 = dict(self.ROW2) |
| final_row2[self.COL0] += " processed" |
| self.assertRowsEqual(final_row0, self._table[0]) |
| self.assertRowsEqual(final_row1, self._table[1]) |
| self.assertRowsEqual(final_row2, self._table[2]) |
