From da567a2679c62a7f35004348ed0ec9e26c93534a Mon Sep 17 00:00:00 2001 From: Misaka_Company Date: Fri, 27 Mar 2026 13:48:07 +0800 Subject: [PATCH] feat(extractor): Add DiscreteMaterialPlanExtractor class for batch data extraction - Create DiscreteMaterialPlanExtractor class orchestrating web operations and post-processing - Implement extract_and_process() combining download + Excel conversion - Implement post_process_downloads() using ExcelConverter for merge - Add extract_from_file() convenience function for file-based ID input - Add component test file verifying class structure and methods - Follow existing patterns: verbose logging, explicit parameters, English comments --- tests/test_extractor_component.py | 120 ++++++++ utils/discrete_material_plan/extractor.py | 334 ++++++++++++++++++++++ 2 files changed, 454 insertions(+) create mode 100644 tests/test_extractor_component.py create mode 100644 utils/discrete_material_plan/extractor.py diff --git a/tests/test_extractor_component.py b/tests/test_extractor_component.py new file mode 100644 index 0000000..924c01d --- /dev/null +++ b/tests/test_extractor_component.py @@ -0,0 +1,120 @@ +""" +Test DiscreteMaterialPlanExtractor component structure +""" + +import sys +from pathlib import Path + +# Add project root to Python path +PROJECT_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(PROJECT_ROOT)) + +# Load environment variables +from dotenv import load_dotenv +load_dotenv(PROJECT_ROOT / ".env") + +# Configure browser path (not used in this test, but consistent with other tests) +import os +os.environ["PLAYWRIGHT_BROWSERS_PATH"] = os.getenv("PLAYWRIGHT_BROWSERS_PATH", "") + +print("=" * 60) +print("Testing DiscreteMaterialPlanExtractor Component") +print("=" * 60) + +try: + # Test 1: Import the component + print("\n[1/4] Importing DiscreteMaterialPlanExtractor...") + from utils.discrete_material_plan.extractor import ( + DiscreteMaterialPlanExtractor, + extract_from_file + ) + print("[OK] Import successful") + + # Test 2: Verify class exists and has expected methods + print("\n[2/4] Verifying class structure...") + expected_methods = [ + '__init__', + '_print', + '_cleanup_temp_files', + '_get_login_url', + '_chunk_order_ids', + '_download_batches', + 'post_process_downloads', + 'extract_and_process', + 'extract_from_file' + ] + + missing_methods = [] + for method in expected_methods: + if not hasattr(DiscreteMaterialPlanExtractor, method): + missing_methods.append(method) + + if missing_methods: + print(f"[FAIL] Missing methods: {missing_methods}") + raise AttributeError(f"Missing methods: {missing_methods}") + + print(f"[OK] All {len(expected_methods)} expected methods found") + + # Test 3: Test instantiation with dummy parameters (no browser) + print("\n[3/4] Testing class instantiation...") + extractor = DiscreteMaterialPlanExtractor( + username="test_user", + password="test_password", + base_url="https://example.com", + headless=True, + ignore_https_errors=True, + verbose=True, + download_dir=None, + batch_size=10 + ) + + # Verify attributes are set correctly + assert extractor.username == "test_user", "Username not set correctly" + assert extractor.password == "test_password", "Password not set correctly" + assert extractor.base_url == "https://example.com", "Base URL not set correctly" + assert extractor.headless == True, "Headless not set correctly" + assert extractor.ignore_https_errors == True, "HTTPS errors setting not set correctly" + assert extractor.verbose == True, "Verbose not set correctly" + assert extractor.batch_size == 10, "Batch size not set correctly" + assert hasattr(extractor, 'download_dir'), "Download dir not set" + + print("[OK] Instantiation successful, all attributes verified") + + # Test 4: Test helper methods + print("\n[4/4] Testing helper methods...") + + # Test _get_login_url + login_url = extractor._get_login_url() + expected_url = "https://example.com/yonbip/resources/uap/rbac/login/main/index.html" + assert login_url == expected_url, f"Login URL incorrect: {login_url}" + print(f"[OK] _get_login_url() returns: {login_url}") + + # Test _chunk_order_ids + test_ids = ["ID1", "ID2", "ID3", "ID4", "ID5"] + extractor.batch_size = 2 + chunks = extractor._chunk_order_ids(test_ids) + expected_chunks = [["ID1", "ID2"], ["ID3", "ID4"], ["ID5"]] + assert chunks == expected_chunks, f"Chunking incorrect: {chunks}" + print(f"[OK] _chunk_order_ids() works correctly: {chunks}") + + # Test _print (should not raise) + extractor._print("Test message") + print("[OK] _print() works correctly") + + # Cleanup + extractor._cleanup_temp_files() + + print("\n" + "=" * 60) + print("[SUCCESS] All component tests passed!") + print("=" * 60) + +except Exception as e: + print(f"\n[ERROR] Component test failed!") + print(f"Error: {e}") + import traceback + traceback.print_exc() + sys.exit(1) + +print("\n" + "=" * 60) +print("Test completed successfully") +print("=" * 60) diff --git a/utils/discrete_material_plan/extractor.py b/utils/discrete_material_plan/extractor.py new file mode 100644 index 0000000..f826fb0 --- /dev/null +++ b/utils/discrete_material_plan/extractor.py @@ -0,0 +1,334 @@ +""" +High-level extractor component for discrete material plan data extraction. +Provides batch processing with progress reporting and error handling. +""" + +import tempfile +import shutil +from pathlib import Path +from typing import List, Optional +from playwright.sync_api import sync_playwright + +from utils.auth import login +from .excel_converter import ExcelConverter +from .extractor_core import ( + navigate_to_discrete_material_page, + setup_query_interface, + execute_batch_download_workflow, +) + + +class DiscreteMaterialPlanExtractor: + """ + High-level extractor for discrete material plan data. + Handles session management, batch processing, and post-processing. + """ + + def __init__( + self, + username: str, + password: str, + base_url: str, + headless: bool = True, + ignore_https_errors: bool = True, + verbose: bool = True, + download_dir: Optional[str] = None, + batch_size: int = 10, + ): + """ + Initialize the extractor. + + Args: + username: Login username + password: Login password + base_url: Base URL for the ERP system + headless: Whether to run browser in headless mode + ignore_https_errors: Whether to ignore HTTPS certificate errors + verbose: Whether to print detailed logs + download_dir: Directory to save downloaded files (default: temp dir) + batch_size: Number of order IDs per batch + """ + self.username = username + self.password = password + self.base_url = base_url + self.headless = headless + self.ignore_https_errors = ignore_https_errors + self.verbose = verbose + self.batch_size = batch_size + + # Set up download directory + if download_dir: + self.download_dir = Path(download_dir) + self.download_dir.mkdir(parents=True, exist_ok=True) + else: + self.download_dir = Path(tempfile.mkdtemp()) + self._is_temp_dir = True + + self._print(f"Download directory: {self.download_dir}") + + def _print(self, *args, **kwargs): + """Print log message if verbose mode is enabled.""" + if self.verbose: + print(*args, **kwargs) + + def _cleanup_temp_files(self): + """Remove temporary files created during extraction.""" + if hasattr(self, "_is_temp_dir") and self._is_temp_dir: + try: + shutil.rmtree(self.download_dir) + self._print(f"Cleaned up temp directory: {self.download_dir}") + except Exception as e: + self._print(f"Warning: Failed to clean up temp directory: {e}") + + def _get_login_url(self) -> str: + """ + Construct the complete login URL. + + Returns: + str: Complete login page URL + """ + base = self.base_url.rstrip("/") + return f"{base}/yonbip/resources/uap/rbac/login/main/index.html" + + def _chunk_order_ids(self, order_ids: List[str]) -> List[List[str]]: + """ + Split order IDs into batches. + + Args: + order_ids: List of order IDs to process + + Returns: + List of batches, where each batch is a list of order IDs + """ + chunks = [] + for i in range(0, len(order_ids), self.batch_size): + chunks.append(order_ids[i : i + self.batch_size]) + return chunks + + def _download_batches( + self, order_ids: List[str] + ) -> List[str]: + """ + Download data for all batches of order IDs. + + Args: + order_ids: List of order IDs to download + + Returns: + List of paths to downloaded files + """ + downloaded_files = [] + chunks = self._chunk_order_ids(order_ids) + + self._print(f"Processing {len(order_ids)} order IDs in {len(chunks)} batches") + + with sync_playwright() as playwright: + # Login + self._print("Logging in...") + browser, context, page, main_frame = login( + playwright=playwright, + username=self.username, + password=self.password, + url=self._get_login_url(), + headless=self.headless, + ignore_https_errors=self.ignore_https_errors, + verbose=self.verbose, + ) + + try: + # Navigate to discrete material page + self._print("Navigating to discrete material plan page...") + work_frame, page1 = navigate_to_discrete_material_page( + main_frame, page + ) + + # Setup query interface + self._print("Setting up query interface...") + setup_query_interface(work_frame) + + # Process each batch + for batch_index, batch in enumerate(chunks): + self._print( + f"Processing batch {batch_index + 1}/{len(chunks)} ({len(batch)} orders)" + ) + + file_path = execute_batch_download_workflow( + work_frame=work_frame, + page=page1, + order_ids=batch, + batch_index=batch_index, + download_dir=str(self.download_dir), + ) + + downloaded_files.append(file_path) + self._print(f"Downloaded: {file_path}") + + finally: + # Cleanup + self._print("Closing browser...") + context.close() + browser.close() + + return downloaded_files + + def post_process_downloads( + self, downloaded_files: List[str], output_file: Optional[str] = None + ) -> Optional[str]: + """ + Convert and merge downloaded Excel files into structured format. + + Args: + downloaded_files: List of paths to downloaded Excel files + output_file: Path to save merged result (default: merged_result.xlsx in download_dir) + + Returns: + Path to the merged output file + """ + if not downloaded_files: + self._print("No files to post-process") + return None + + self._print(f"Post-processing {len(downloaded_files)} downloaded files...") + + converter = ExcelConverter(verbose=self.verbose) + all_dfs = [] + + # Convert each file + for i, file_path in enumerate(downloaded_files): + self._print(f"Converting file {i + 1}/{len(downloaded_files)}: {file_path}") + df = converter.convert(input_file=file_path) # type: ignore + all_dfs.append(df) + + # Merge all DataFrames + if all_dfs: + import pandas as pd + + merged_df = pd.concat(all_dfs, ignore_index=True) + + # Determine output file path + if output_file: + output_path = Path(output_file) + else: + output_path = self.download_dir / "merged_result.xlsx" + + # Save merged result + merged_df.to_excel(output_path, index=False) + self._print(f"Merged result saved to: {output_path}") + self._print(f"Total rows: {len(merged_df)}") + + return str(output_path) + + return None + + def extract_and_process( + self, order_ids: List[str], output_file: Optional[str] = None + ) -> Optional[str]: + """ + Complete extraction workflow: download + post-process. + + Args: + order_ids: List of order IDs to extract + output_file: Path to save merged result (optional) + + Returns: + Path to the final merged output file + """ + # Download all batches + downloaded_files = self._download_batches(order_ids) + + # Post-process: convert and merge + output_path = self.post_process_downloads(downloaded_files, output_file) + + return output_path + + def extract_from_file( + self, id_file: str, output_file: Optional[str] = None + ) -> Optional[str]: + """ + Extract data from order IDs listed in a file. + + Args: + id_file: Path to file containing order IDs (one per line) + output_file: Path to save merged result (optional) + + Returns: + Path to the final merged output file + """ + # Read order IDs from file + id_path = Path(id_file) + if not id_path.exists(): + raise FileNotFoundError(f"ID file not found: {id_file}") + + with open(id_path, "r", encoding="utf-8") as f: + order_ids = [line.strip() for line in f if line.strip()] + + self._print(f"Loaded {len(order_ids)} order IDs from {id_file}") + + # Extract and process + return self.extract_and_process(order_ids, output_file) + + +def extract_from_file( + id_file: str, + output_file: Optional[str] = None, + username: Optional[str] = None, + password: Optional[str] = None, + base_url: Optional[str] = None, + headless: bool = True, + ignore_https_errors: bool = True, + verbose: bool = True, + download_dir: Optional[str] = None, + batch_size: int = 10, +) -> Optional[str]: + """ + Convenience function to extract data from order IDs in a file. + + Args: + id_file: Path to file containing order IDs (one per line) + output_file: Path to save merged result (optional) + username: Login username (required if not in env) + password: Login password (required if not in env) + base_url: Base URL for ERP system (required if not in env) + headless: Whether to run browser in headless mode + ignore_https_errors: Whether to ignore HTTPS certificate errors + verbose: Whether to print detailed logs + download_dir: Directory to save downloaded files (optional) + batch_size: Number of order IDs per batch + + Returns: + Path to the final merged output file + + Note: + If username, password, or base_url are not provided, + they will be read from environment variables. + """ + import os + + # Get credentials from parameters or environment + if not username: + username = os.getenv("ERP_USERNAME") + if not password: + password = os.getenv("ERP_PASSWORD") + if not base_url: + base_url = os.getenv("ERP_URL") + + if not username or not password or not base_url: + raise ValueError( + "username, password, and base_url must be provided either as parameters or environment variables" + ) + + extractor = DiscreteMaterialPlanExtractor( + username=username, + password=password, + base_url=base_url, + headless=headless, + ignore_https_errors=ignore_https_errors, + verbose=verbose, + download_dir=download_dir, + batch_size=batch_size, + ) + + try: + return extractor.extract_from_file(id_file, output_file) + finally: + extractor._cleanup_temp_files()