tests: fix tests

This commit is contained in:
Bryce 2023-05-20 12:43:32 -07:00 committed by Bryce Drennan
parent 39dffa9166
commit 38ac0b7f54
3 changed files with 9 additions and 3 deletions

View File

@ -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):

Binary file not shown.

After

Width:  |  Height:  |  Size: 351 KiB

View File

@ -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,