From d2f9ae691cc78a6b322f81cb2dd290795808e77e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Magalh=C3=A3es?= <joamag@gmail.com>
Date: Tue, 1 Oct 2024 12:51:31 +0100
Subject: [PATCH] chore: support for multiple account selection

---
 CHANGELOG.md               |  2 +-
 src/weby_pilot/__main__.py |  2 +-
 src/weby_pilot/bpi.py      | 30 +++++++++++++++++++++++++-----
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3e18743..c3e31d9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
-*
+* Support for account selection
 
 ### Changed
 
diff --git a/src/weby_pilot/__main__.py b/src/weby_pilot/__main__.py
index b1c5f8c..eaee058 100644
--- a/src/weby_pilot/__main__.py
+++ b/src/weby_pilot/__main__.py
@@ -3,7 +3,7 @@
 
 from .bpi import BpiAPI
 
-print(BpiAPI().download_account_report(report_indexes=range(0, 2)))
+print(BpiAPI().download_account_report(report_indexes=range(0, 2), account_index=1))
 
 # BpiAPI().download_report(section="Extrato Investimento", report_indexes=range(0, 8))
 # BpiAPI().download_invoice()
diff --git a/src/weby_pilot/bpi.py b/src/weby_pilot/bpi.py
index 91b69c8..8df09bc 100644
--- a/src/weby_pilot/bpi.py
+++ b/src/weby_pilot/bpi.py
@@ -2,12 +2,14 @@
 # -*- coding: utf-8 -*-
 
 from enum import Enum
+from time import sleep
 from os import environ
 from os.path import basename
 from datetime import datetime
-from typing import IO, Literal, Sequence, Tuple, cast
+from typing import IO, Literal, Sequence, cast
 from selenium.webdriver.common.by import By
 from selenium.webdriver.common.keys import Keys
+from selenium.webdriver.support.ui import Select
 
 from .base import WebyAPI
 from .common import FileType
@@ -119,12 +121,14 @@ def download_report(
         self,
         section: ReportSections = "Extrato Conta",
         report_indexes: Sequence[int] = (0,),
+        account_index: int = 0,
     ) -> Sequence["BpiDocument"]:
         docs: list[BpiDocument] = []
         with self.driver_ctx():
             self.login()
             self.select_section("Consultas")
             self.select_side_menu(cast(BpiSideSections, section))
+            self.select_account(account_index)
             for report_index in report_indexes:
                 self.click_extract(row_index=report_index)
                 docs.append(
@@ -142,17 +146,21 @@ def download_report(
         return docs
 
     def download_account_report(
-        self, report_indexes: Sequence[int] = (0,)
+        self, report_indexes: Sequence[int] = (0,), account_index: int = 0
     ) -> Sequence["BpiDocument"]:
         return self.download_report(
-            section="Extrato Conta", report_indexes=report_indexes
+            section="Extrato Conta",
+            report_indexes=report_indexes,
+            account_index=account_index,
         )
 
     def download_investing_report(
-        self, report_indexes: Sequence[int] = (0,)
+        self, report_indexes: Sequence[int] = (0,), account_index: int = 0
     ) -> Sequence["BpiDocument"]:
         return self.download_report(
-            section="Extrato Investimento", report_indexes=report_indexes
+            section="Extrato Investimento",
+            report_indexes=report_indexes,
+            account_index=account_index,
         )
 
     def download_card_report(self, card_index=0, report_indexes: Sequence[int] = (0,)):
@@ -199,6 +207,18 @@ def select_filters(self, date_range: SelectDateRange, document_type: DocumentTyp
         filter = self.get_element(By.XPATH, "//*[@value='Filtrar']")
         filter.click()
 
+    def select_account(self, index: int = 0, timeout=5.0):
+        account_element = self.get_element(
+            By.XPATH, "//div[text()='Conta']/following-sibling::div/*/select"
+        )
+        account_select = Select(account_element)
+        selected_index = account_select.options.index(
+            account_select.first_selected_option
+        )
+        if selected_index != index:
+            account_select.select_by_index(index)
+            sleep(timeout)
+
     def click_extract(self, row_index=0, wait_download: bool = True):
         open_extract = self.get_element(
             By.XPATH,