Fix bugs surfaced by tests. Also remove extraneous import