Add more typing annotations to test_engine.py