Fix the broken test because of the dask api-change.

[SchoolGuy]

Fixes: #666

[philippjfr]

Thanks, looks good to me.

[jbednar]

Thanks! I guess we don't need to add any version pinning or version checks, because we're only testing with the latest Dask, and this code is only in our tests, not our actual codebase?

[SchoolGuy]

I can't say because my goal was just to fix the test. Problem is/was when setting up the dev-environment for me that my PyCharm complained about a lot of missing packages which sometimes resulted in tests failing which are actually not failing. If you like I open another PR for that. (@philippjfr @jbednar )