diff --git a/internal_filesystem/builtin/apps/com.micropythonos.osupdate/assets/osupdate.py b/internal_filesystem/builtin/apps/com.micropythonos.osupdate/assets/osupdate.py index deceb590..82236fa3 100644 --- a/internal_filesystem/builtin/apps/com.micropythonos.osupdate/assets/osupdate.py +++ b/internal_filesystem/builtin/apps/com.micropythonos.osupdate/assets/osupdate.py @@ -2,10 +2,9 @@ import lvgl as lv import requests import ujson import time -import _thread from mpos.apps import Activity -from mpos import PackageManager, ConnectivityManager +from mpos import PackageManager, ConnectivityManager, TaskManager, DownloadManager import mpos.info import mpos.ui @@ -256,11 +255,9 @@ class OSUpdate(Activity): self.progress_bar.align(lv.ALIGN.BOTTOM_MID, 0, -50) self.progress_bar.set_range(0, 100) self.progress_bar.set_value(0, False) - try: - _thread.stack_size(mpos.apps.good_stack_size()) - _thread.start_new_thread(self.update_with_lvgl, (self.download_update_url,)) - except Exception as e: - print("Could not start update_with_lvgl thread: ", e) + + # Use TaskManager instead of _thread for async download + TaskManager.create_task(self.perform_update()) def force_update_clicked(self): if self.download_update_url and (self.force_update.get_state() & lv.STATE.CHECKED): @@ -275,33 +272,36 @@ class OSUpdate(Activity): self.set_state(UpdateState.CHECKING_UPDATE) self.schedule_show_update_info() - def progress_callback(self, percent): + async def async_progress_callback(self, percent): + """Async progress callback for DownloadManager.""" print(f"OTA Update: {percent:.1f}%") - self.update_ui_threadsafe_if_foreground(self.progress_bar.set_value, int(percent), True) - self.update_ui_threadsafe_if_foreground(self.progress_label.set_text, f"OTA Update: {percent:.2f}%") - time.sleep_ms(100) + # UI updates are safe from async context in MicroPythonOS (runs on main thread) + if self.has_foreground(): + self.progress_bar.set_value(int(percent), True) + self.progress_label.set_text(f"OTA Update: {percent:.2f}%") + await TaskManager.sleep_ms(50) - # Custom OTA update with LVGL progress - def update_with_lvgl(self, url): - """Download and install update in background thread. + async def perform_update(self): + """Download and install update using async patterns. Supports automatic pause/resume on wifi loss. """ + url = self.download_update_url + try: # Loop to handle pause/resume cycles while self.has_foreground(): - # Use UpdateDownloader to handle the download - result = self.update_downloader.download_and_install( + # Use UpdateDownloader to handle the download (now async) + result = await self.update_downloader.download_and_install( url, - progress_callback=self.progress_callback, + progress_callback=self.async_progress_callback, should_continue_callback=self.has_foreground ) if result['success']: # Update succeeded - set boot partition and restart - self.update_ui_threadsafe_if_foreground(self.status_label.set_text,"Update finished! Restarting...") - # Small delay to show the message - time.sleep(5) + self.status_label.set_text("Update finished! Restarting...") + await TaskManager.sleep(5) self.update_downloader.set_boot_partition_and_restart() return @@ -314,8 +314,7 @@ class OSUpdate(Activity): print(f"OSUpdate: Download paused at {percent:.1f}% ({bytes_written}/{total_size} bytes)") self.set_state(UpdateState.DOWNLOAD_PAUSED) - # Wait for wifi to return - # ConnectivityManager will notify us via callback when network returns + # Wait for wifi to return using async sleep print("OSUpdate: Waiting for network to return...") check_interval = 2 # Check every 2 seconds max_wait = 300 # 5 minutes timeout @@ -324,19 +323,19 @@ class OSUpdate(Activity): while elapsed < max_wait and self.has_foreground(): if self.connectivity_manager.is_online(): print("OSUpdate: Network reconnected, waiting for stabilization...") - time.sleep(2) # Let routing table and DNS fully stabilize + await TaskManager.sleep(2) # Let routing table and DNS fully stabilize print("OSUpdate: Resuming download") self.set_state(UpdateState.DOWNLOADING) break # Exit wait loop and retry download - time.sleep(check_interval) + await TaskManager.sleep(check_interval) elapsed += check_interval if elapsed >= max_wait: # Timeout waiting for network msg = f"Network timeout during download.\n{bytes_written}/{total_size} bytes written.\nPress 'Update OS' to retry." - self.update_ui_threadsafe_if_foreground(self.status_label.set_text, msg) - self.update_ui_threadsafe_if_foreground(self.install_button.remove_state, lv.STATE.DISABLED) + self.status_label.set_text(msg) + self.install_button.remove_state(lv.STATE.DISABLED) self.set_state(UpdateState.ERROR) return @@ -344,32 +343,40 @@ class OSUpdate(Activity): else: # Update failed with error (not pause) - error_msg = result.get('error', 'Unknown error') - bytes_written = result.get('bytes_written', 0) - total_size = result.get('total_size', 0) - - if "cancelled" in error_msg.lower(): - msg = ("Update cancelled by user.\n\n" - f"{bytes_written}/{total_size} bytes downloaded.\n" - "Press 'Update OS' to resume.") - else: - # Use friendly error message - friendly_msg = self._get_user_friendly_error(Exception(error_msg)) - progress_info = f"\n\nProgress: {bytes_written}/{total_size} bytes" - if bytes_written > 0: - progress_info += "\n\nPress 'Update OS' to resume." - msg = friendly_msg + progress_info - - self.set_state(UpdateState.ERROR) - self.update_ui_threadsafe_if_foreground(self.status_label.set_text, msg) - self.update_ui_threadsafe_if_foreground(self.install_button.remove_state, lv.STATE.DISABLED) # allow retry + self._handle_update_error(result) return except Exception as e: - msg = self._get_user_friendly_error(e) + "\n\nPress 'Update OS' to retry." - self.set_state(UpdateState.ERROR) - self.update_ui_threadsafe_if_foreground(self.status_label.set_text, msg) - self.update_ui_threadsafe_if_foreground(self.install_button.remove_state, lv.STATE.DISABLED) # allow retry + self._handle_update_exception(e) + + def _handle_update_error(self, result): + """Handle update error result - extracted for DRY.""" + error_msg = result.get('error', 'Unknown error') + bytes_written = result.get('bytes_written', 0) + total_size = result.get('total_size', 0) + + if "cancelled" in error_msg.lower(): + msg = ("Update cancelled by user.\n\n" + f"{bytes_written}/{total_size} bytes downloaded.\n" + "Press 'Update OS' to resume.") + else: + # Use friendly error message + friendly_msg = self._get_user_friendly_error(Exception(error_msg)) + progress_info = f"\n\nProgress: {bytes_written}/{total_size} bytes" + if bytes_written > 0: + progress_info += "\n\nPress 'Update OS' to resume." + msg = friendly_msg + progress_info + + self.set_state(UpdateState.ERROR) + self.status_label.set_text(msg) + self.install_button.remove_state(lv.STATE.DISABLED) # allow retry + + def _handle_update_exception(self, e): + """Handle update exception - extracted for DRY.""" + msg = self._get_user_friendly_error(e) + "\n\nPress 'Update OS' to retry." + self.set_state(UpdateState.ERROR) + self.status_label.set_text(msg) + self.install_button.remove_state(lv.STATE.DISABLED) # allow retry # Business Logic Classes: @@ -386,19 +393,22 @@ class UpdateState: ERROR = "error" class UpdateDownloader: - """Handles downloading and installing OS updates.""" + """Handles downloading and installing OS updates using async DownloadManager.""" - def __init__(self, requests_module=None, partition_module=None, connectivity_manager=None): + # Chunk size for partition writes (must be 4096 for ESP32 flash) + CHUNK_SIZE = 4096 + + def __init__(self, partition_module=None, connectivity_manager=None, download_manager=None): """Initialize with optional dependency injection for testing. Args: - requests_module: HTTP requests module (defaults to requests) partition_module: ESP32 Partition module (defaults to esp32.Partition if available) connectivity_manager: ConnectivityManager instance for checking network during download + download_manager: DownloadManager module for async downloads (defaults to mpos.DownloadManager) """ - self.requests = requests_module if requests_module else requests self.partition_module = partition_module self.connectivity_manager = connectivity_manager + self.download_manager = download_manager # For testing injection self.simulate = False # Download state for pause/resume @@ -406,6 +416,13 @@ class UpdateDownloader: self.bytes_written_so_far = 0 self.total_size_expected = 0 + # Internal state for chunk processing + self._current_partition = None + self._block_index = 0 + self._chunk_buffer = b'' + self._should_continue = True + self._progress_callback = None + # Try to import Partition if not provided if self.partition_module is None: try: @@ -442,14 +459,87 @@ class UpdateDownloader: return any(indicator in error_str or indicator in error_repr for indicator in network_indicators) - def download_and_install(self, url, progress_callback=None, should_continue_callback=None): - """Download firmware and install to OTA partition. + def _setup_partition(self): + """Initialize the OTA partition for writing.""" + if not self.simulate and self._current_partition is None: + current = self.partition_module(self.partition_module.RUNNING) + self._current_partition = current.get_next_update() + print(f"UpdateDownloader: Writing to partition: {self._current_partition}") + + async def _process_chunk(self, chunk): + """Process a downloaded chunk - buffer and write to partition. + + Note: Progress reporting is handled by DownloadManager, not here. + This method only handles buffering and writing to partition. + + Args: + chunk: bytes data received from download + """ + # Check if we should continue (user cancelled) + if not self._should_continue: + return + + # Check network connection + if self.connectivity_manager: + is_online = self.connectivity_manager.is_online() + elif ConnectivityManager._instance: + is_online = ConnectivityManager._instance.is_online() + else: + is_online = True + + if not is_online: + print("UpdateDownloader: Network lost during chunk processing") + self.is_paused = True + raise OSError(-113, "Network lost during download") + + # Track total bytes received + self._total_bytes_received += len(chunk) + + # Add chunk to buffer + self._chunk_buffer += chunk + + # Write complete 4096-byte blocks + while len(self._chunk_buffer) >= self.CHUNK_SIZE: + block = self._chunk_buffer[:self.CHUNK_SIZE] + self._chunk_buffer = self._chunk_buffer[self.CHUNK_SIZE:] + + if not self.simulate: + self._current_partition.writeblocks(self._block_index, block) + + self._block_index += 1 + self.bytes_written_so_far += len(block) + + # Note: Progress is reported by DownloadManager via progress_callback parameter + # We don't calculate progress here to avoid duplicate/incorrect progress updates + + async def _flush_buffer(self): + """Flush remaining buffer with padding to complete the download.""" + if self._chunk_buffer: + # Pad the last chunk to 4096 bytes + remaining = len(self._chunk_buffer) + padded = self._chunk_buffer + b'\xFF' * (self.CHUNK_SIZE - remaining) + print(f"UpdateDownloader: Padding final chunk from {remaining} to {self.CHUNK_SIZE} bytes") + + if not self.simulate: + self._current_partition.writeblocks(self._block_index, padded) + + self.bytes_written_so_far += self.CHUNK_SIZE + self._chunk_buffer = b'' + + # Final progress update + if self._progress_callback and self.total_size_expected > 0: + percent = (self.bytes_written_so_far / self.total_size_expected) * 100 + await self._progress_callback(min(percent, 100.0)) + + async def download_and_install(self, url, progress_callback=None, should_continue_callback=None): + """Download firmware and install to OTA partition using async DownloadManager. Supports pause/resume on wifi loss using HTTP Range headers. Args: url: URL to download firmware from - progress_callback: Optional callback function(percent: float) + progress_callback: Optional async callback function(percent: float) + Called by DownloadManager with progress 0-100 should_continue_callback: Optional callback function() -> bool Returns False to cancel download @@ -460,9 +550,6 @@ class UpdateDownloader: - 'total_size': int - 'error': str (if success=False) - 'paused': bool (if paused due to wifi loss) - - Raises: - Exception: If download or installation fails """ result = { 'success': False, @@ -472,135 +559,99 @@ class UpdateDownloader: 'paused': False } - try: - # Get OTA partition - next_partition = None - if not self.simulate: - current = self.partition_module(self.partition_module.RUNNING) - next_partition = current.get_next_update() - print(f"UpdateDownloader: Writing to partition: {next_partition}") + # Store callbacks for use in _process_chunk + self._progress_callback = progress_callback + self._should_continue = True + self._total_bytes_received = 0 - # Start download (or resume if we have bytes_written_so_far) - headers = {} + try: + # Setup partition + self._setup_partition() + + # Initialize block index from resume position + self._block_index = self.bytes_written_so_far // self.CHUNK_SIZE + + # Build headers for resume + headers = None if self.bytes_written_so_far > 0: - headers['Range'] = f'bytes={self.bytes_written_so_far}-' + headers = {'Range': f'bytes={self.bytes_written_so_far}-'} print(f"UpdateDownloader: Resuming from byte {self.bytes_written_so_far}") - response = self.requests.get(url, stream=True, headers=headers) + # Get the download manager (use injected one for testing, or global) + dm = self.download_manager if self.download_manager else DownloadManager - # For initial download, get total size + # Create wrapper for chunk callback that checks should_continue + async def chunk_handler(chunk): + if should_continue_callback and not should_continue_callback(): + self._should_continue = False + raise Exception("Download cancelled by user") + await self._process_chunk(chunk) + + # For initial download, we need to get total size first + # DownloadManager doesn't expose Content-Length directly, so we estimate if self.bytes_written_so_far == 0: - total_size = int(response.headers.get('Content-Length', 0)) - result['total_size'] = round_up_to_multiple(total_size, 4096) - self.total_size_expected = result['total_size'] + # We'll update total_size_expected as we download + # For now, set a placeholder that will be updated + self.total_size_expected = 0 + + # Download with streaming chunk callback + # Progress is reported by DownloadManager via progress_callback + print(f"UpdateDownloader: Starting async download from {url}") + success = await dm.download_url( + url, + chunk_callback=chunk_handler, + progress_callback=progress_callback, # Let DownloadManager handle progress + headers=headers + ) + + if success: + # Flush any remaining buffered data + await self._flush_buffer() + + result['success'] = True + result['bytes_written'] = self.bytes_written_so_far + result['total_size'] = self.bytes_written_so_far # Actual size downloaded + + # Final 100% progress callback + if self._progress_callback: + await self._progress_callback(100.0) + + # Reset state for next download + self.is_paused = False + self.bytes_written_so_far = 0 + self.total_size_expected = 0 + self._current_partition = None + self._block_index = 0 + self._chunk_buffer = b'' + self._total_bytes_received = 0 + + print(f"UpdateDownloader: Download complete ({result['bytes_written']} bytes)") else: - # For resume, use the stored total size - # (Content-Length will be the remaining bytes, not total) + # Download failed but not due to exception + result['error'] = "Download failed" + result['bytes_written'] = self.bytes_written_so_far result['total_size'] = self.total_size_expected - print(f"UpdateDownloader: Download target {result['total_size']} bytes") - - chunk_size = 4096 - bytes_written = self.bytes_written_so_far - block_index = bytes_written // chunk_size - - while True: - # Check if we should continue (user cancelled) - if should_continue_callback and not should_continue_callback(): - result['error'] = "Download cancelled by user" - response.close() - return result - - # Check network connection before reading - if self.connectivity_manager: - is_online = self.connectivity_manager.is_online() - elif ConnectivityManager._instance: - is_online = ConnectivityManager._instance.is_online() - else: - is_online = True - - if not is_online: - print("UpdateDownloader: Network lost (pre-check), pausing download") - self.is_paused = True - self.bytes_written_so_far = bytes_written - result['paused'] = True - result['bytes_written'] = bytes_written - response.close() - return result - - # Read next chunk (may raise exception if network drops) - try: - chunk = response.raw.read(chunk_size) - except Exception as read_error: - # Check if this is a network error that should trigger pause - if self._is_network_error(read_error): - print(f"UpdateDownloader: Network error during read ({read_error}), pausing") - self.is_paused = True - self.bytes_written_so_far = bytes_written - result['paused'] = True - result['bytes_written'] = bytes_written - try: - response.close() - except: - pass - return result - else: - # Non-network error, re-raise - raise - - if not chunk: - break - - # Pad last chunk if needed - if len(chunk) < chunk_size: - print(f"UpdateDownloader: Padding chunk {block_index} from {len(chunk)} to {chunk_size} bytes") - chunk = chunk + b'\xFF' * (chunk_size - len(chunk)) - - # Write to partition - if not self.simulate: - next_partition.writeblocks(block_index, chunk) - - bytes_written += len(chunk) - self.bytes_written_so_far = bytes_written - block_index += 1 - - # Update progress - if progress_callback and result['total_size'] > 0: - percent = (bytes_written / result['total_size']) * 100 - progress_callback(percent) - - # Small delay to avoid hogging CPU - time.sleep_ms(100) - - response.close() - result['bytes_written'] = bytes_written - - # Check if complete - if bytes_written >= result['total_size']: - result['success'] = True - self.is_paused = False - self.bytes_written_so_far = 0 # Reset for next download - self.total_size_expected = 0 - print(f"UpdateDownloader: Download complete ({bytes_written} bytes)") - else: - result['error'] = f"Incomplete download: {bytes_written} < {result['total_size']}" - print(f"UpdateDownloader: {result['error']}") - except Exception as e: + error_msg = str(e) + + # Check if cancelled by user + if "cancelled" in error_msg.lower(): + result['error'] = error_msg + result['bytes_written'] = self.bytes_written_so_far + result['total_size'] = self.total_size_expected # Check if this is a network error that should trigger pause - if self._is_network_error(e): + elif self._is_network_error(e): print(f"UpdateDownloader: Network error ({e}), pausing download") self.is_paused = True - # Only update bytes_written_so_far if we actually wrote bytes in this attempt - # Otherwise preserve the existing state (important for resume failures) - if result.get('bytes_written', 0) > 0: - self.bytes_written_so_far = result['bytes_written'] result['paused'] = True result['bytes_written'] = self.bytes_written_so_far - result['total_size'] = self.total_size_expected # Preserve total size for UI + result['total_size'] = self.total_size_expected else: # Non-network error - result['error'] = str(e) + result['error'] = error_msg + result['bytes_written'] = self.bytes_written_so_far + result['total_size'] = self.total_size_expected print(f"UpdateDownloader: Error during download: {e}") return result diff --git a/tests/network_test_helper.py b/tests/network_test_helper.py index c811c1fe..05349c5e 100644 --- a/tests/network_test_helper.py +++ b/tests/network_test_helper.py @@ -680,3 +680,217 @@ class MockTime: def clear_sleep_calls(self): """Clear the sleep call history.""" self._sleep_calls = [] + + +class MockDownloadManager: + """ + Mock DownloadManager for testing async downloads. + + Simulates the mpos.DownloadManager module for testing without actual network I/O. + Supports chunk_callback mode for streaming downloads. + """ + + def __init__(self): + """Initialize mock download manager.""" + self.download_data = b'' + self.should_fail = False + self.fail_after_bytes = None + self.headers_received = None + self.url_received = None + self.call_history = [] + self.chunk_size = 1024 # Default chunk size for streaming + + async def download_url(self, url, outfile=None, total_size=None, + progress_callback=None, chunk_callback=None, headers=None): + """ + Mock async download with flexible output modes. + + Simulates the real DownloadManager behavior including: + - Streaming chunks via chunk_callback + - Progress reporting via progress_callback (based on total size) + - Network failure simulation + + Args: + url: URL to download + outfile: Path to write file (optional) + total_size: Expected size for progress tracking (optional) + progress_callback: Async callback for progress updates (optional) + chunk_callback: Async callback for streaming chunks (optional) + headers: HTTP headers dict (optional) + + Returns: + bytes: Downloaded content (if outfile and chunk_callback are None) + bool: True if successful (when using outfile or chunk_callback) + """ + self.url_received = url + self.headers_received = headers + + # Record call in history + self.call_history.append({ + 'url': url, + 'outfile': outfile, + 'total_size': total_size, + 'headers': headers, + 'has_progress_callback': progress_callback is not None, + 'has_chunk_callback': chunk_callback is not None + }) + + if self.should_fail: + if outfile or chunk_callback: + return False + return None + + # Check for immediate failure (fail_after_bytes=0) + if self.fail_after_bytes is not None and self.fail_after_bytes == 0: + raise OSError(-113, "ECONNABORTED") + + # Stream data in chunks + bytes_sent = 0 + chunks = [] + total_data_size = len(self.download_data) + + # Use provided total_size or actual data size for progress calculation + effective_total_size = total_size if total_size else total_data_size + + while bytes_sent < total_data_size: + # Check if we should simulate network failure + if self.fail_after_bytes is not None and bytes_sent >= self.fail_after_bytes: + raise OSError(-113, "ECONNABORTED") + + chunk = self.download_data[bytes_sent:bytes_sent + self.chunk_size] + + if chunk_callback: + await chunk_callback(chunk) + elif outfile: + # For file mode, we'd write to file (mock just tracks) + pass + else: + chunks.append(chunk) + + bytes_sent += len(chunk) + + # Report progress (like real DownloadManager does) + if progress_callback and effective_total_size > 0: + percent = round((bytes_sent * 100) / effective_total_size) + await progress_callback(percent) + + # Return based on mode + if outfile or chunk_callback: + return True + else: + return b''.join(chunks) + + def set_download_data(self, data): + """ + Configure the data to return from downloads. + + Args: + data: Bytes to return from download + """ + self.download_data = data + + def set_should_fail(self, should_fail): + """ + Configure whether downloads should fail. + + Args: + should_fail: True to make downloads fail + """ + self.should_fail = should_fail + + def set_fail_after_bytes(self, bytes_count): + """ + Configure network failure after specified bytes. + + Args: + bytes_count: Number of bytes to send before failing + """ + self.fail_after_bytes = bytes_count + + def clear_history(self): + """Clear the call history.""" + self.call_history = [] + + +class MockTaskManager: + """ + Mock TaskManager for testing async operations. + + Provides mock implementations of TaskManager methods for testing. + """ + + def __init__(self): + """Initialize mock task manager.""" + self.tasks_created = [] + self.sleep_calls = [] + + @classmethod + def create_task(cls, coroutine): + """ + Mock create_task - just runs the coroutine synchronously for testing. + + Args: + coroutine: Coroutine to execute + + Returns: + The coroutine (for compatibility) + """ + # In tests, we typically run with asyncio.run() so just return the coroutine + return coroutine + + @staticmethod + async def sleep(seconds): + """ + Mock async sleep. + + Args: + seconds: Number of seconds to sleep (ignored in mock) + """ + pass # Don't actually sleep in tests + + @staticmethod + async def sleep_ms(milliseconds): + """ + Mock async sleep in milliseconds. + + Args: + milliseconds: Number of milliseconds to sleep (ignored in mock) + """ + pass # Don't actually sleep in tests + + @staticmethod + async def wait_for(awaitable, timeout): + """ + Mock wait_for with timeout. + + Args: + awaitable: Coroutine to await + timeout: Timeout in seconds (ignored in mock) + + Returns: + Result of the awaitable + """ + return await awaitable + + @staticmethod + def notify_event(): + """ + Create a mock async event. + + Returns: + A simple mock event object + """ + class MockEvent: + def __init__(self): + self._set = False + + async def wait(self): + pass + + def set(self): + self._set = True + + def is_set(self): + return self._set + + return MockEvent() diff --git a/tests/test_osupdate.py b/tests/test_osupdate.py index 16e52fd2..88687edd 100644 --- a/tests/test_osupdate.py +++ b/tests/test_osupdate.py @@ -1,12 +1,13 @@ import unittest import sys +import asyncio # Add parent directory to path so we can import network_test_helper # When running from unittest.sh, we're in internal_filesystem/, so tests/ is ../tests/ sys.path.insert(0, '../tests') # Import network test helpers -from network_test_helper import MockNetwork, MockRequests, MockJSON +from network_test_helper import MockNetwork, MockRequests, MockJSON, MockDownloadManager class MockPartition: @@ -42,6 +43,11 @@ sys.path.append('builtin/apps/com.micropythonos.osupdate/assets') from osupdate import UpdateChecker, UpdateDownloader, round_up_to_multiple +def run_async(coro): + """Helper to run async coroutines in sync tests.""" + return asyncio.get_event_loop().run_until_complete(coro) + + class TestUpdateChecker(unittest.TestCase): """Test UpdateChecker class.""" @@ -218,38 +224,37 @@ class TestUpdateChecker(unittest.TestCase): class TestUpdateDownloader(unittest.TestCase): - """Test UpdateDownloader class.""" + """Test UpdateDownloader class with async DownloadManager.""" def setUp(self): - self.mock_requests = MockRequests() + self.mock_download_manager = MockDownloadManager() self.mock_partition = MockPartition self.downloader = UpdateDownloader( - requests_module=self.mock_requests, - partition_module=self.mock_partition + partition_module=self.mock_partition, + download_manager=self.mock_download_manager ) def test_download_and_install_success(self): """Test successful download and install.""" # Create 8KB of test data (2 blocks of 4096 bytes) test_data = b'A' * 8192 - self.mock_requests.set_next_response( - status_code=200, - headers={'Content-Length': '8192'}, - content=test_data - ) + self.mock_download_manager.set_download_data(test_data) + self.mock_download_manager.chunk_size = 4096 progress_calls = [] - def progress_cb(percent): + async def progress_cb(percent): progress_calls.append(percent) - result = self.downloader.download_and_install( - "http://example.com/update.bin", - progress_callback=progress_cb - ) + async def run_test(): + return await self.downloader.download_and_install( + "http://example.com/update.bin", + progress_callback=progress_cb + ) + + result = run_async(run_test()) self.assertTrue(result['success']) self.assertEqual(result['bytes_written'], 8192) - self.assertEqual(result['total_size'], 8192) self.assertIsNone(result['error']) # MicroPython unittest doesn't have assertGreater self.assertTrue(len(progress_calls) > 0, "Should have progress callbacks") @@ -257,21 +262,21 @@ class TestUpdateDownloader(unittest.TestCase): def test_download_and_install_cancelled(self): """Test cancelled download.""" test_data = b'A' * 8192 - self.mock_requests.set_next_response( - status_code=200, - headers={'Content-Length': '8192'}, - content=test_data - ) + self.mock_download_manager.set_download_data(test_data) + self.mock_download_manager.chunk_size = 4096 call_count = [0] def should_continue(): call_count[0] += 1 return call_count[0] < 2 # Cancel after first chunk - result = self.downloader.download_and_install( - "http://example.com/update.bin", - should_continue_callback=should_continue - ) + async def run_test(): + return await self.downloader.download_and_install( + "http://example.com/update.bin", + should_continue_callback=should_continue + ) + + result = run_async(run_test()) self.assertFalse(result['success']) self.assertIn("cancelled", result['error'].lower()) @@ -280,44 +285,46 @@ class TestUpdateDownloader(unittest.TestCase): """Test that last chunk is properly padded.""" # 5000 bytes - not a multiple of 4096 test_data = b'B' * 5000 - self.mock_requests.set_next_response( - status_code=200, - headers={'Content-Length': '5000'}, - content=test_data - ) + self.mock_download_manager.set_download_data(test_data) + self.mock_download_manager.chunk_size = 4096 - result = self.downloader.download_and_install( - "http://example.com/update.bin" - ) + async def run_test(): + return await self.downloader.download_and_install( + "http://example.com/update.bin" + ) + + result = run_async(run_test()) self.assertTrue(result['success']) - # Should be rounded up to 8192 (2 * 4096) - self.assertEqual(result['total_size'], 8192) + # Should be padded to 8192 (2 * 4096) + self.assertEqual(result['bytes_written'], 8192) def test_download_with_network_error(self): """Test download with network error during transfer.""" - self.mock_requests.set_exception(Exception("Network error")) + self.mock_download_manager.set_should_fail(True) - result = self.downloader.download_and_install( - "http://example.com/update.bin" - ) + async def run_test(): + return await self.downloader.download_and_install( + "http://example.com/update.bin" + ) + + result = run_async(run_test()) self.assertFalse(result['success']) self.assertIsNotNone(result['error']) - self.assertIn("Network error", result['error']) def test_download_with_zero_content_length(self): """Test download with missing or zero Content-Length.""" test_data = b'C' * 1000 - self.mock_requests.set_next_response( - status_code=200, - headers={}, # No Content-Length header - content=test_data - ) + self.mock_download_manager.set_download_data(test_data) + self.mock_download_manager.chunk_size = 1000 - result = self.downloader.download_and_install( - "http://example.com/update.bin" - ) + async def run_test(): + return await self.downloader.download_and_install( + "http://example.com/update.bin" + ) + + result = run_async(run_test()) # Should still work, just with unknown total size initially self.assertTrue(result['success']) @@ -325,60 +332,58 @@ class TestUpdateDownloader(unittest.TestCase): def test_download_progress_callback_called(self): """Test that progress callback is called during download.""" test_data = b'D' * 8192 - self.mock_requests.set_next_response( - status_code=200, - headers={'Content-Length': '8192'}, - content=test_data - ) + self.mock_download_manager.set_download_data(test_data) + self.mock_download_manager.chunk_size = 4096 progress_values = [] - def track_progress(percent): + async def track_progress(percent): progress_values.append(percent) - result = self.downloader.download_and_install( - "http://example.com/update.bin", - progress_callback=track_progress - ) + async def run_test(): + return await self.downloader.download_and_install( + "http://example.com/update.bin", + progress_callback=track_progress + ) + + result = run_async(run_test()) self.assertTrue(result['success']) # Should have at least 2 progress updates (for 2 chunks of 4096) self.assertTrue(len(progress_values) >= 2) # Last progress should be 100% - self.assertEqual(progress_values[-1], 100.0) + self.assertEqual(progress_values[-1], 100) def test_download_small_file(self): """Test downloading a file smaller than one chunk.""" test_data = b'E' * 100 # Only 100 bytes - self.mock_requests.set_next_response( - status_code=200, - headers={'Content-Length': '100'}, - content=test_data - ) + self.mock_download_manager.set_download_data(test_data) + self.mock_download_manager.chunk_size = 100 - result = self.downloader.download_and_install( - "http://example.com/update.bin" - ) + async def run_test(): + return await self.downloader.download_and_install( + "http://example.com/update.bin" + ) + + result = run_async(run_test()) self.assertTrue(result['success']) # Should be padded to 4096 - self.assertEqual(result['total_size'], 4096) self.assertEqual(result['bytes_written'], 4096) def test_download_exact_chunk_multiple(self): """Test downloading exactly 2 chunks (no padding needed).""" test_data = b'F' * 8192 # Exactly 2 * 4096 - self.mock_requests.set_next_response( - status_code=200, - headers={'Content-Length': '8192'}, - content=test_data - ) + self.mock_download_manager.set_download_data(test_data) + self.mock_download_manager.chunk_size = 4096 - result = self.downloader.download_and_install( - "http://example.com/update.bin" - ) + async def run_test(): + return await self.downloader.download_and_install( + "http://example.com/update.bin" + ) + + result = run_async(run_test()) self.assertTrue(result['success']) - self.assertEqual(result['total_size'], 8192) self.assertEqual(result['bytes_written'], 8192) def test_network_error_detection_econnaborted(self): @@ -417,16 +422,16 @@ class TestUpdateDownloader(unittest.TestCase): """Test that download pauses when network error occurs during read.""" # Set up mock to raise network error after first chunk test_data = b'G' * 16384 # 4 chunks - self.mock_requests.set_next_response( - status_code=200, - headers={'Content-Length': '16384'}, - content=test_data, - fail_after_bytes=4096 # Fail after first chunk - ) + self.mock_download_manager.set_download_data(test_data) + self.mock_download_manager.chunk_size = 4096 + self.mock_download_manager.set_fail_after_bytes(4096) # Fail after first chunk - result = self.downloader.download_and_install( - "http://example.com/update.bin" - ) + async def run_test(): + return await self.downloader.download_and_install( + "http://example.com/update.bin" + ) + + result = run_async(run_test()) self.assertFalse(result['success']) self.assertTrue(result['paused']) @@ -436,29 +441,27 @@ class TestUpdateDownloader(unittest.TestCase): def test_download_resumes_from_saved_position(self): """Test that download resumes from the last written position.""" # Simulate partial download - test_data = b'H' * 12288 # 3 chunks self.downloader.bytes_written_so_far = 8192 # Already downloaded 2 chunks self.downloader.total_size_expected = 12288 - # Server should receive Range header + # Server should receive Range header - only remaining data remaining_data = b'H' * 4096 # Last chunk - self.mock_requests.set_next_response( - status_code=206, # Partial content - headers={'Content-Length': '4096'}, # Remaining bytes - content=remaining_data - ) + self.mock_download_manager.set_download_data(remaining_data) + self.mock_download_manager.chunk_size = 4096 - result = self.downloader.download_and_install( - "http://example.com/update.bin" - ) + async def run_test(): + return await self.downloader.download_and_install( + "http://example.com/update.bin" + ) + + result = run_async(run_test()) self.assertTrue(result['success']) self.assertEqual(result['bytes_written'], 12288) # Check that Range header was set - last_request = self.mock_requests.last_request - self.assertIsNotNone(last_request) - self.assertIn('Range', last_request['headers']) - self.assertEqual(last_request['headers']['Range'], 'bytes=8192-') + self.assertIsNotNone(self.mock_download_manager.headers_received) + self.assertIn('Range', self.mock_download_manager.headers_received) + self.assertEqual(self.mock_download_manager.headers_received['Range'], 'bytes=8192-') def test_resume_failure_preserves_state(self): """Test that resume failures preserve download state for retry.""" @@ -466,12 +469,16 @@ class TestUpdateDownloader(unittest.TestCase): self.downloader.bytes_written_so_far = 245760 # 60 chunks already downloaded self.downloader.total_size_expected = 3391488 - # Resume attempt fails immediately with EHOSTUNREACH (network not ready) - self.mock_requests.set_exception(OSError(-118, "EHOSTUNREACH")) + # Resume attempt fails immediately with network error + self.mock_download_manager.set_download_data(b'') + self.mock_download_manager.set_fail_after_bytes(0) # Fail immediately - result = self.downloader.download_and_install( - "http://example.com/update.bin" - ) + async def run_test(): + return await self.downloader.download_and_install( + "http://example.com/update.bin" + ) + + result = run_async(run_test()) # Should pause, not fail self.assertFalse(result['success'])