mirror of
https://github.com/brycedrennan/imaginAIry
synced 2024-10-31 03:20:40 +00:00
tests: fix tests
This commit is contained in:
parent
39dffa9166
commit
38ac0b7f54
@ -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 |
@ -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,
|
||||
|
Loading…
Reference in New Issue
Block a user