Add tests for methods within scripts of `trainer`

[jsingh811]

Description

Add a pytest script to test functions inside trainer.

Acceptance Criteria

-[ ] Add unit tests for functions in trainer