tests: fix tests

imaginairy/utils/model_cache.py

@@ -171,7 +171,7 @@ class GPUModelCache:
             or self.gpu_cache.memory_usage + bytes_to_free
             > get_mem_free_total(self.device)
         ):
-            msg = f"Unable to make {bytes_to_free} space on GPU. \n{self.stats_msg()}"
+            msg = f"Unable to make {bytes_to_free / (1024 ** 2):.1f} MB space on {self.device}. \n{self.stats_msg()}"
             raise RuntimeError(msg)
 
     def make_cpu_space(self, bytes_to_free):

tests/test_cmds.py

@@ -1,3 +1,5 @@
+from unittest import mock
+
 from click.testing import CliRunner
 
 from imaginairy import ImaginePrompt, LazyLoadingImage, surprise_me
@@ -6,10 +8,12 @@ from imaginairy.cli.edit_demo import edit_demo_cmd
 from imaginairy.cli.imagine import imagine_cmd
 from imaginairy.cli.main import aimg
 from imaginairy.cli.upscale import upscale_cmd
+from imaginairy.utils.model_cache import GPUModelCache
 from tests import TESTS_FOLDER
 
 
-def test_imagine_cmd():
+def test_imagine_cmd(monkeypatch):
+    monkeypatch.setattr(GPUModelCache, "make_gpu_space", mock.MagicMock())
     runner = CliRunner()
     result = runner.invoke(
         imagine_cmd,
@@ -30,7 +34,8 @@ def test_imagine_cmd():
     assert result.exit_code == 0
 
 
-def test_edit_cmd():
+def test_edit_cmd(monkeypatch):
+    monkeypatch.setattr(GPUModelCache, "make_gpu_space", mock.MagicMock())
     runner = CliRunner()
     result = runner.invoke(
         edit_cmd,
@@ -74,6 +79,7 @@ def test_edit_demo(monkeypatch):
         ]
 
     monkeypatch.setattr(surprise_me, "surprise_me_prompts", mock_surprise_me_prompts)
+    monkeypatch.setattr(GPUModelCache, "make_gpu_space", mock.MagicMock())
     surprise_me.generic_prompts = []
     result = runner.invoke(
         edit_demo_cmd,