Refactor & Tests

[dejanmiletic89]

Problem

Describe the purpose of this change. What problem is being solved and why?

Solution

Describe the approach you took. Link to any relevant bugs, issues, docs, or other resources.

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update
• Infrastructure change (CI configs, etc)
• Non-code change (docs, etc)
• None of the above: (explain here)

Test Plan

Describe specific steps for validating this change.