28 lines
798 B
Python
28 lines
798 B
Python
# tests/test_gate.py
|
|
import io
|
|
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from gates import gate
|
|
|
|
def test_route_tuple_places_image_at_chosen():
|
|
B = object()
|
|
t = gate.route_tuple(2, "IMG", B, max_routes=5)
|
|
assert t == (B, B, "IMG", B, B)
|
|
|
|
def test_route_tuple_length_is_max():
|
|
B = object()
|
|
assert len(gate.route_tuple(0, "IMG", B, max_routes=10)) == 10
|
|
|
|
def test_mask_from_stash_none_is_zeros():
|
|
img = torch.zeros((1, 6, 4, 3))
|
|
m = gate.mask_from_stash(None, img)
|
|
assert m.shape == (1, 6, 4) and float(m.max()) == 0.0
|
|
|
|
def test_mask_from_stash_decodes_png():
|
|
buf = io.BytesIO(); Image.new("L", (4, 6), 255).save(buf, "PNG")
|
|
img = torch.zeros((1, 6, 4, 3))
|
|
m = gate.mask_from_stash(buf.getvalue(), img)
|
|
assert m.shape == (1, 6, 4) and float(m.min()) > 0.99
|