Unit test for distillation model handlers

# Feature
## Desired Behavior / Functionality
We need to have unit tests in place for the distillation model handler.

## What Needs to Be Done
Write unit tests testing all the functionality.

## How Can It Be Tested
- Unit tests passing