Added additional google sheets functions

abr-testing/abr_testing/automation/google_sheets_tool.py

@@ -48,20 +48,21 @@ def open_worksheet(self, tab_number: int) -> Any:
         """Open individual worksheet within a googlesheet."""
         return self.spread_sheet.get_worksheet(tab_number)
 
-    def create_worksheet(self, tab_name: int) -> None:
+    def create_worksheet(self, title: str) -> None:
         """Create a worksheet with tab name. Existing spreadsheet needed."""
         try:
-            self.spread_sheet.add_worksheet(tab_name, rows="1000", cols="26")
+            new_sheet = self.spread_sheet.add_worksheet(title, rows="2000", cols="26")
+            return new_sheet.id
         except gspread.exceptions.APIError:
-            print("Work Sheet already exists")
+            print("Sheet already exists.")
 
     def write_header(self, header: List) -> None:
         """Write Header to first row if not present."""
         header_list = self.worksheet.row_values(1)
         if header_list != header:
             self.worksheet.insert_row(header, self.row_index)
 
-    def write_to_row(self, data: List) -> None:
+    def write_to_row(self, data: List, title: str = "Sheet1") -> None:
         """Write data into a row in a List[] format."""
         try:
             self.row_index += 1
@@ -106,11 +107,24 @@ def batch_delete_rows(self, row_indices: List[int]) -> None:
         }
         self.spread_sheet.batch_update(body=delete_body)
 
+    def batch_update_cells(
+        self, sheet_title: str, data: List[List[str]], start_column: str, start_row: int
+    ) -> None:
+        """Writes to multiple cells at once in a specific sheet."""
+        sheet = self.spread_sheet.worksheet(sheet_title)
+        for idx, values in enumerate(data):
+            column = chr(ord(start_column) + idx)  # Convert index to column letter
+            location = f"{column}{start_row}:{column}{start_row + len(values) - 1}"
+            cells_to_update = sheet.range(location)
+            for cell, value in zip(cells_to_update, values):
+                cell.value = value
+            sheet.update_cells(cells_to_update)
+
     def update_cell(
-        self, row: int, column: int, single_data: Any
+        self, sheet_title: str, row: int, column: int, single_data: Any
     ) -> Tuple[int, int, Any]:
         """Update ONE individual cell according to a row and column."""
-        self.worksheet.update_cell(row, column, single_data)
+        self.spread_sheet.worksheet(sheet_title).update_cell(row, column, single_data)
         return row, column, single_data
 
     def get_all_data(self) -> List[Dict[str, Any]]:
@@ -121,6 +135,15 @@ def get_column(self, column_number: int) -> Set[str]:
         """Get all values in column."""
         return self.worksheet.col_values(column_number)
 
+    def get_cell(self, cell: str) -> Any:
+        """Get cell value with location ex: 'A1'."""
+        return self.worksheet.acell(cell).value
+
+    def get_single_col_range(self, range: str) -> List:
+        """Get cell values from one column range."""
+        values_range = self.worksheet.range(range)
+        return [cell.value for cell in values_range]
+
     def get_index_row(self) -> int:
         """Check for the next available row to write too."""
         row_index = len(self.get_column(1))
@@ -166,6 +189,9 @@ def create_line_chart(
         titles: List[str],
         series: List[Dict[str, Any]],
         domains: List[Dict[str, Any]],
+        axis: Dict[str, Any],
+        col_position: int = 0,
+        sheet_id: str = "0",
     ) -> None:
         """Create chart of data on google sheet."""
         request_body = {
@@ -178,10 +204,7 @@ def create_line_chart(
                                 "basicChart": {
                                     "chartType": "LINE",
                                     "legendPosition": "RIGHT_LEGEND",
-                                    "axis": [
-                                        {"position": "BOTTOM_AXIS", "title": titles[1]},
-                                        {"position": "LEFT_AXIS", "title": titles[2]},
-                                    ],
+                                    "axis": axis,
                                     "domains": domains,
                                     "series": series,
                                     "headerCount": 1,
@@ -190,9 +213,9 @@ def create_line_chart(
                             "position": {
                                 "overlayPosition": {
                                     "anchorCell": {
-                                        "sheetId": 0,
+                                        "sheetId": sheet_id,
                                         "rowIndex": 1,
-                                        "columnIndex": 1,
+                                        "columnIndex": col_position,
                                     }
                                 }
                             },