I have a mock DatasetWriter object and a function to write an array into the object.

def write_dataset(dataset: DatasetWriter, image: np.ndarray):
    dataset.write(image)

Now, I want to write a unittest for this function.

def test_write_dataset():
    output_dataset = Mock(DatasetWriter)
    write_dataset(dataset=output_dataset, image=np.array([[1, 2, 3], [4, 5, 6]]))
    print(output_dataset.read())
    assert (output_dataset.read() == np.array([[1, 2, 3], [4, 5, 6]])).all()

The output:

============================================================= FAILURES ==============================================================
________________________________________________________ test_write_dataset _________________________________________________________

    def test_write_dataset():
        output_dataset = Mock(DatasetWriter)
        write_dataset(dataset=output_dataset, image=np.array([[1, 2, 3], [4, 5, 6]]))
        print(output_dataset.read())
>       assert (output_dataset.read() == np.array([[1, 2, 3], [4, 5, 6]])).all()
E       AssertionError: assert False
E        +  where False = <built-in method all of numpy.ndarray object at 0x7f21fdba24b0>()
E        +    where <built-in method all of numpy.ndarray object at 0x7f21fdba24b0> = <Mock name='mock.read()' id='139783967501328'> == array([[1, 2, 3],\n       [4, 5, 6]]).all
E        +      where <Mock name='mock.read()' id='139783967501328'> = <Mock name='mock.read' id='139783967501088'>()
E        +        where <Mock name='mock.read' id='139783967501088'> = <Mock spec='DatasetWriter' id='139783967500752'>.read
E        +      and   array([[1, 2, 3],\n       [4, 5, 6]]) = <built-in function array>([[1, 2, 3], [4, 5, 6]])
E        +        where <built-in function array> = np.array

tests/liveeo/test_processing.py:134: AssertionError
------------------------------------------------------- Captured stdout call --------------------------------------------------------
<Mock name='mock.read()' id='139783967501328'>
========================================================= 1 failed in 0.55s =========================================================

My question is that the mock object output_dataset can not perform the operation of write in the write_dataset function. How could I write a unittest for the function write_dataset. Thanks.