[autotest] Use AfeHostInfoStore to obtain host information in test_push

This is not really needed to remove autoserv dependencies on the AFE,
but this allows us to remove the duplicated functions in site_utils.

BUG=chromium:678430
TEST=test_push

Change-Id: Id725e97b6d726fdf68c063582d45c640d4541bbe
Reviewed-on: https://chromium-review.googlesource.com/442656
Commit-Ready: Prathmesh Prabhu <pprabhu@chromium.org>
Tested-by: Prathmesh Prabhu <pprabhu@chromium.org>
Reviewed-by: Shuqian Zhao <shuqianz@chromium.org>
diff --git a/site_utils/test_push.py b/site_utils/test_push.py
index 2118c9e..19f8036 100755
--- a/site_utils/test_push.py
+++ b/site_utils/test_push.py
@@ -46,6 +46,7 @@
 from autotest_lib.server import utils
 from autotest_lib.server.cros import provision
 from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
+from autotest_lib.server.hosts import afe_store
 from autotest_lib.site_utils import gmail_lib
 from autotest_lib.site_utils.suite_scheduler import constants
 
@@ -380,11 +381,12 @@
             for job_id in job_ids]
     hostnames = set([hqe.host.hostname for hqe in hqes])
     for hostname in hostnames:
-        found_build = site_utils.get_build_from_afe(hostname, AFE)
-        if found_build != build:
+        host_info_store = afe_store.AfeStore(hostname, AFE)
+        info = host_info_store.get()
+        if info.build != build:
             raise TestPushException('DUT is not imaged properly. Host %s has '
                                     'build %s, while build %s is expected.' %
-                                    (hostname, found_build, build))
+                                    (hostname, info.build, build))
 
 
 def test_suite(suite_name, expected_results, arguments, use_shard=False,