Add more checks on elapsed time in `test_wait.go`

[ryan-gang]

Currently we don't check the time elapsed when the WAIT command is supposed to return early. But, as highlighted by @andy1li, this might lead to cases where user never implements WAIT properly, just sleeps for time_duration, and returns.

This can be mitigated by adding checks on elapsed_time in cases where WAIT returns early.

[linear[bot]]

CC-1055 Add more checks on elapsed time in `test_wait.go`