Move tests to tests.rs file