From 16c3defe20b8736519d095b8f4187fc629c0ab37 Mon Sep 17 00:00:00 2001
From: HappySix <33394488+FredWuCZ@users.noreply.github.com>
Date: Thu, 20 Jun 2024 19:03:13 +0800
Subject: [PATCH] Change resolution before saving snapshot in VirtualBox (#47)

* Initailize aws support

* Add README for the VM server

* Refactor OSWorld for supporting more cloud services.

* Initialize vmware and aws implementation v1, waiting for verification

* Initlize files for azure, gcp and virtualbox support

* Debug on the VMware provider

* Fix on aws interface mapping

* Fix instance type

* Refactor

* Clean

* Add Azure provider

* hk region; debug

* Fix lock

* Remove print

* Remove key_name requirements when allocating aws vm

* Clean README

* Fix reset

* Fix bugs

* Add VirtualBox and Azure providers

* Add VirtualBox OVF link

* Raise exception on macOS host

* Init RAEDME for VBox

* Update VirtualBox VM download link

* Update requirements and setup.py; Improve robustness on Windows

* Fix network adapter

* Go through on Windows machine

* Add default adapter option

* Fix minor error

* Change resolution before creating snapshot

* Fix small error

* Change default provider option

---------

Co-authored-by: Timothyxxx <384084775@qq.com>
Co-authored-by: XinyuanWangCS <xywang626@gmail.com>
Co-authored-by: Tianbao Xie <47296835+Timothyxxx@users.noreply.github.com>
---
 desktop_env/providers/virtualbox/manager.py  | 18 +++++++++++++++---
 desktop_env/providers/virtualbox/provider.py |  1 -
 main.py                                      |  2 +-
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/desktop_env/providers/virtualbox/manager.py b/desktop_env/providers/virtualbox/manager.py
index 4f9e569..9d1b546 100644
--- a/desktop_env/providers/virtualbox/manager.py
+++ b/desktop_env/providers/virtualbox/manager.py
@@ -228,7 +228,7 @@ def _install_vm(vm_name, vms_dir, downloaded_file_name, original_vm_name="Ubuntu
     if not start_vm(vm_name):
         raise ValueError("Error encountered during installation, please rerun the code for retrying.")
 
-    def get_vm_ip(vm_name, max_retries=20):
+    def get_vm_ip(vm_name):
         command = f'VBoxManage guestproperty get "{vm_name}" /VirtualBox/GuestInfo/Net/0/V4/IP'
         result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8")
         if result.returncode == 0:
@@ -236,9 +236,17 @@ def _install_vm(vm_name, vms_dir, downloaded_file_name, original_vm_name="Ubuntu
         else:
             logger.error(f"Get VM IP failed: {result.stderr}")
             return None
+        
+    def change_resolution(vm_name, resolution=(1920, 1080, 32)):
+        command = f'VBoxManage controlvm "{vm_name}" setvideomodehint {" ".join(map(str, resolution))}'
+        result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8")
+        if result.returncode == 0:
+            return True
+        else:
+            return False
 
     # Function used to check whether the virtual machine is ready
-    def download_screenshot():
+    def download_screenshot(vm_name):
         ip = get_vm_ip(vm_name)
         url = f"http://{ip}:5000/screenshot"
         try:
@@ -253,10 +261,14 @@ def _install_vm(vm_name, vms_dir, downloaded_file_name, original_vm_name="Ubuntu
         return False
 
     # Try downloading the screenshot until successful
-    while not download_screenshot():
+    while not download_screenshot(vm_name):
         logger.info("Check whether the virtual machine is ready...")
         time.sleep(RETRY_INTERVAL)
 
+    if not change_resolution(vm_name):
+        logger.error(f"Change resolution failed.")
+        raise
+
     logger.info("Virtual machine is ready. Start to make a snapshot on the virtual machine. It would take a while...")
 
     def create_vm_snapshot(vm_name, max_retries=20):
diff --git a/desktop_env/providers/virtualbox/provider.py b/desktop_env/providers/virtualbox/provider.py
index 71e086e..81cda08 100644
--- a/desktop_env/providers/virtualbox/provider.py
+++ b/desktop_env/providers/virtualbox/provider.py
@@ -56,7 +56,6 @@ class VirtualBoxProvider(Provider):
             
 
     def start_emulator(self, path_to_vm: str, headless: bool):
-        print("Starting VirtualBox VM...")
         logger.info("Starting VirtualBox VM...")
 
         while True:
diff --git a/main.py b/main.py
index bc87d54..4c6817f 100644
--- a/main.py
+++ b/main.py
@@ -49,7 +49,7 @@ def human_agent():
     parser = argparse.ArgumentParser()
     parser.add_argument('-p', '--path', type=str, default="", help="Path to the virtual machine.")
     parser.add_argument('-e', '--example', type=str, help="Path to the example json file.")
-    parser.add_argument('-s', '--snapshot', type=str, help="Name of the snapshot to load.")
+    parser.add_argument('-s', '--snapshot', type=str, default="init_state", help="Name of the snapshot to load.")
     parser.add_argument('-r', '--region', type=str, help="(For VirtualBox) Name of the bridged adapter. (For AWS) Name of the region.")
     args = parser.parse_args(sys.argv[1:])