#!/usr/bin/env python

from numpy import all, array, dstack

from .. import image


def test_get_channels():
    x = array([[1, 2, 3], [2, 3, 4], [3, 4, 5]])
    y = dstack((x, x, x))
    for z in image.get_channels(y):
        assert all(z == x)


def test_map_to_channels():
    def f(x):
        return x + 1
    x = array([[1, 2, 3], [2, 3, 4], [3, 4, 5]])
    y = dstack((x, x, x))
    z = image.map_to_channels(
        f,
        lambda shape: shape,
        y,
    )
    assert all(z == y + 1)


if __name__ == '__main__':
    pass

@@ -38,16 +38,26 @@
 ort all,
+ allclose,
  array, 
@@ -62,16 +62,69 @@
 , dstack
+%0Afrom os import remove%0Afrom os.path import sep, split
 %0A%0Afrom .
@@ -137,16 +137,23 @@
 rt image
+, utils
 %0A%0A%0Adef t
@@ -567,16 +567,291 @@
  + 1)%0A%0A%0A
+def test_read_save():%0A    f = split(__file__)%5B0%5D + sep + 'drink.png'%0A    x = image.read(f)%0A    tmp = utils.rand_filename(f)%0A    try:%0A        image.save(x, tmp)%0A        y = image.read(tmp)%0A        assert allclose(x, y, rtol=1e-4, atol=1e-1)%0A    finally:%0A        remove(tmp)%0A%0A%0A
 if __nam

Add more unit tests for the image module.