Increasing test coverage and fixing bugs found while writing tests

Replace with list comprehension?

```python
features = [dict(zip(feature_names, row)) for row in X]
```