refactoring osworld constructor

agential-ai · Jan 6, 2025 · bfaeadc · bfaeadc
1 parent 087bfcd
commit bfaeadc
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/agential/benchmarks/computer_use/osworld/osworld.py b/agential/benchmarks/computer_use/osworld/osworld.py
@@ -88,20 +88,19 @@ def __init__(
                       and the parent `BaseComputerUseBenchmark` class.
         """
         super().__init__(**kwargs)
-        self.ubuntu0_vmx = ubuntu0_vmx
 
         DesktopEnv.__init__ = initializer
 
-        ubuntu0_vmx = kwargs.get("path_to_vm")
+        self.ubuntu0_vmx = kwargs.get("path_to_vm")
 
         try:
             self.env = DesktopEnv(**kwargs)
         except:
             try:
-                vmrun_command = ['vmrun', 'start', ubuntu0_vmx]
+                vmrun_command = ['vmrun', 'start', self.ubuntu0_vmx]
                 subprocess.run(vmrun_command, check=True)
 
-                self.env = DesktopEnv(path_to_vm=ubuntu0_vmx, **kwargs)
+                self.env = DesktopEnv(path_to_vm=self.ubuntu0_vmx, **kwargs)
 
                 print("VM started successfully.")
             except subprocess.CalledProcessError as e:

diff --git a/tests/benchmarks/computer_use/osworld/test_osworld.py b/tests/benchmarks/computer_use/osworld/test_osworld.py
@@ -14,9 +14,6 @@ def test_init_() -> None:
     Test the __init__ function of OSWorld in the virtual machine.
     """
     env = MagicMock(spec=OSWorld)
-    env.vmware_vm_data = "to_vmware_vm_data_folder"
-    env.ubuntu0 = "to_ubuntu0_folder"
-    env.ubuntu0_vmx = "to_ubuntu0_vmx_file"
 
     assert env.vmware_vm_data == "to_vmware_vm_data_folder"
     assert env.ubuntu0 == "to_ubuntu0_folder"