Clean code; Refactor environment to pass screenshot content instead of path

desktop_env/controllers/setup.py

@@ -2,6 +2,7 @@ import json
 import logging
 import os
 import os.path
+import shutil
 import sqlite3
 import tempfile
 import time
@@ -11,7 +12,6 @@ from datetime import datetime, timedelta
 from typing import Any, Union, Optional
 from typing import Dict, List
 
-import shutil
 import requests
 from playwright.sync_api import sync_playwright, TimeoutError
 from pydrive.auth import GoogleAuth
@@ -25,6 +25,7 @@ logger = logging.getLogger("desktopenv.setup")
 
 FILE_PATH = os.path.dirname(os.path.abspath(__file__))
 
+
 class SetupController:
     def __init__(self, vm_ip: str, cache_dir: str):
         self.vm_ip: str = vm_ip
@@ -60,39 +61,6 @@ class SetupController:
 
             logger.info("SETUP: %s(%s)", setup_function, str(parameters))
 
-        # self._download_setup(config)
-        # self._change_wallpaper(config)
-        # self._tidy_desktop(config) todo: implement this
-        # self._open_setup(config)
-        # can add other setup steps
-
-    # ZDY_COMMENT: merged with launch
-    # def _command_setup(self, command: str):
-    # """
-    # Directly send a command into the virtual machine os for setting up.
-    # """
-    # payload = json.dumps({"command": command})
-    # headers = {
-    # 'Content-Type': 'application/json'
-    # }
-    # timeout = 5
-    # timout_whitelist = ["vlc"]
-    #
-    # try:
-    #
-    # response = requests.post(self.http_server + "/execute", headers=headers, data=payload, timeout=timeout)
-    # if response.status_code == 200:
-    # print("Command executed successfully:", response.text)
-    # else:
-    # print("Failed to execute command. Status code:", response.status_code)
-    # except requests.exceptions.Timeout as e:
-    # if command in timout_whitelist:
-    # print("Command executed successfully:", command)
-    # else:
-    # print("An error occurred while trying to execute the command:", e)
-    # except requests.exceptions.RequestException as e:
-    # print("An error occurred while trying to execute the command:", e)
-
     def _download_setup(self, files: List[Dict[str, str]]):
         """
         Args:
@@ -140,11 +108,6 @@ class SetupController:
                 if not downloaded:
                     raise requests.RequestException(f"Failed to download {url}. No retries left. Error: {e}")
 
-            # payload = json.dumps({"url": url, "path": path})
-            # headers = {
-            # 'Content-Type': 'application/json'
-            # }
-
             form = MultipartEncoder({
                 "file_path": path,
                 "file_data": (os.path.basename(path), open(cache_path, "rb"))
@@ -163,6 +126,41 @@ class SetupController:
             except requests.exceptions.RequestException as e:
                 logger.error("An error occurred while trying to send the request: %s", e)
 
+    def _upload_file_setup(self, files: List[Dict[str, str]]):
+        """
+        Args:
+            files (List[Dict[str, str]]): files to download. lisf of dict like
+              {
+                "local_path": str, the local path to the file to upload
+                "path": str, the path on the VM to store the downloaded file
+              }
+        """
+        for f in files:
+            local_path: str = f["local_path"]
+            path: str = f["path"]
+
+            if not os.path.exists(local_path):
+                logger.error(f"Setup Upload - Invalid local path ({local_path}).")
+                return
+
+            form = MultipartEncoder({
+                "file_path": path,
+                "file_data": (os.path.basename(path), open(local_path, "rb"))
+            })
+            headers = {"Content-Type": form.content_type}
+            logger.debug(form.content_type)
+
+            # send request to server to upload file
+            try:
+                logger.debug("REQUEST ADDRESS: %s", self.http_server + "/setup" + "/upload")
+                response = requests.post(self.http_server + "/setup" + "/upload", headers=headers, data=form)
+                if response.status_code == 200:
+                    logger.info("Command executed successfully: %s", response.text)
+                else:
+                    logger.error("Failed to upload file. Status code: %s", response.text)
+            except requests.exceptions.RequestException as e:
+                logger.error("An error occurred while trying to send the request: %s", e)
+
     def _change_wallpaper_setup(self, path: str):
         # if not config:
         # return
@@ -559,7 +557,7 @@ class SetupController:
                 try:
                     page.goto(url, timeout=60000)
                 except:
-                    logger.warning("Opening %s exceeds time limit", url) # only for human test
+                    logger.warning("Opening %s exceeds time limit", url)  # only for human test
                 logger.info(f"Opened new page: {url}")
                 settings = json.load(open(config['settings_file']))
                 email, password = settings['email'], settings['password']