Improve tests for inference

[davanstrien]

At the moment there are tests for inference but they aren't very pretty. This issue is a meta issue to track progress on making the tests nicer.

Insperation/reading

• read some of the tests for hugging face transformers
• read https://www.integralist.co.uk/posts/mocking-in-python

integration tests

unit tests

• refactor inference code to make it easier to test separate logic #516