From 505ef91bc945c2999f756a57bec74ef652c56dd2 Mon Sep 17 00:00:00 2001
From: Andrew Gunnerson <accounts+github@chiller3.com>
Date: Sat, 11 Feb 2023 20:37:47 -0500
Subject: [PATCH] Use unittest.mock.patch for replacing classes and variables

The mechanism is still the same (modifies global state), but the code is
a bit cleaner.

Signed-off-by: Andrew Gunnerson <accounts+github@chiller3.com>
---
 avbroot/ota.py    | 11 +++--------
 avbroot/vbmeta.py | 10 +++-------
 2 files changed, 6 insertions(+), 15 deletions(-)

diff --git a/avbroot/ota.py b/avbroot/ota.py
index c7418ba..aa54161 100644
--- a/avbroot/ota.py
+++ b/avbroot/ota.py
@@ -10,6 +10,7 @@ import os
 import struct
 import sys
 import subprocess
+import unittest.mock
 import zipfile
 
 # Silence undesired warning
@@ -758,12 +759,6 @@ def match_android_zip64_limit():
     [1] https://cs.android.com/android/platform/superproject/+/android-13.0.0_r18:system/libziparchive/zip_archive.cc;l=692
     '''
 
-    orig_limit = zipfile.ZIP64_LIMIT
-
-    try:
-        # Because Python uses > and Android uses >= 0xffffffff
-        zipfile.ZIP64_LIMIT = 0xfffffffe
-
+    # Because Python uses > and Android uses >= 0xffffffff
+    with unittest.mock.patch('zipfile.ZIP64_LIMIT', 0xfffffffe):
         yield
-    finally:
-        zipfile.ZIP64_LIMIT = orig_limit
diff --git a/avbroot/vbmeta.py b/avbroot/vbmeta.py
index 4f92eac..1015620 100644
--- a/avbroot/vbmeta.py
+++ b/avbroot/vbmeta.py
@@ -1,4 +1,5 @@
 import contextlib
+import unittest.mock
 
 import avbtool
 
@@ -26,14 +27,9 @@ def smuggle_descriptors():
     * call encode on each descriptor
     '''
 
-    orig_kernel = avbtool.AvbKernelCmdlineDescriptor
-
-    avbtool.AvbKernelCmdlineDescriptor = SmuggledViaKernelCmdlineDescriptor
-
-    try:
+    with unittest.mock.patch('avbtool.AvbKernelCmdlineDescriptor',
+                             SmuggledViaKernelCmdlineDescriptor):
         yield
-    finally:
-        avbtool.AvbKernelCmdlineDescriptor = orig_kernel
 
 
 def _get_descriptor_overrides(avb, paths):