commit ffeb89893888594f27a29fc1891f08c08d698d7d
author Alan Zhang <alanjz@google.com> Fri Jul 29 00:39:00 2016 +0000
committer Gerrit Code Review <noreply-gerritcodereview@google.com> Fri Jul 29 00:39:01 2016 +0000
tree e29f19d0ddfad9c23260e9091021d4c8b5205c4d
parent d7f19e942fcea39293d5c5fb331ee26850711a2e
parent 053cfcb536db3312ebde1a79f712a71ea3f4631d

Merge "Adjusted fix for crashes from empty cache_warming_hosts with tests."

cache_warming/cache_warming.py

diff --git a/cache_warming/cache_warming.py b/cache_warming/cache_warming.py
index 5ef444c..986bcd7 100644
--- a/cache_warming/cache_warming.py
+++ b/cache_warming/cache_warming.py
@@ -56,7 +56,14 @@
   if os.path.isfile(HOSTS_JSON_PATH):
     with open(HOSTS_JSON_PATH, 'r') as hosts_json:
       global hit_log
-      hit_log = json.load(hosts_json)
+      try:
+        hit_log = json.load(hosts_json)
+      except ValueError as e:
+        if verbose:
+          print 'Failed to open %s: %s.' % (HOSTS_JSON_PATH, e)
+      finally:
+        if not isinstance(hit_log, dict):
+          hit_log = {}
 
 
 def process_query(qry):
@@ -171,13 +178,7 @@
   sys.stderr = os.fdopen(2, 'w', 1)
   args = set_args()
   verbose = args.verbose
-  try:
-    load_hosts()
-  except ValueError as e:
-    if verbose:
-      print 'Failed to open %s: %s.' % (HOSTS_JSON_PATH, e)
-    save_hosts(hit_log)
-
+  load_hosts()
   server_address = UDP_SERVER_PATH
   try:
     os.remove(server_address)

cache_warming/cache_warming_test.py

diff --git a/cache_warming/cache_warming_test.py b/cache_warming/cache_warming_test.py
index 220edc1..ceed91b 100644
--- a/cache_warming/cache_warming_test.py
+++ b/cache_warming/cache_warming_test.py
@@ -1,6 +1,7 @@
 #!/usr/bin/python
 """Tests for cache_warming.py."""
 
+import os
 import cache_warming
 from wvtest import wvtest
 
@@ -72,5 +73,55 @@
   wvtest.WVPASSEQ(actual, expected)
 
 
+@wvtest.wvtest
+def test_no_cache_warming_hosts():
+  crash = False
+  cache_warming.HOSTS_JSON_PATH = '/tmp/cache_warming_hosts.json'
+  if os.path.exists('/tmp/cache_warming_hosts.json'):
+    os.remove('/tmp/cache_warming_hosts.json')
+  try:
+    cache_warming.load_hosts()
+    cache_warming.warm_cache(53, None)
+  except ValueError:
+    crash = True
+  wvtest.WVFAIL(crash)
+
+
+@wvtest.wvtest
+def test_empty_cache_warming_hosts():
+  crash = False
+  cache_warming.HOSTS_JSON_PATH = '/tmp/cache_warming_hosts.json'
+  if os.path.exists('/tmp/cache_warming_hosts.json'):
+    os.remove('/tmp/cache_warming_hosts.json')
+  open('/tmp/cache_warming_hosts.json', 'w').close()
+  try:
+    cache_warming.load_hosts()
+    cache_warming.warm_cache(53, None)
+  except ValueError:
+    crash = True
+  finally:
+    os.remove('/tmp/cache_warming_hosts.json')
+  wvtest.WVFAIL(crash)
+
+
+@wvtest.wvtest
+def test_wrong_cache_warming_hosts():
+  crash = False
+  cache_warming.HOSTS_JSON_PATH = '/tmp/cache_warming_hosts.json'
+  if os.path.exists('/tmp/cache_warming_hosts.json'):
+    os.remove('/tmp/cache_warming_hosts.json')
+  f = open('/tmp/cache_warming_hosts.json', 'w')
+  f.write('[]')
+  f.close()
+  try:
+    cache_warming.load_hosts()
+    cache_warming.warm_cache(53, None)
+  except ValueError:
+    crash = True
+  finally:
+    os.remove('/tmp/cache_warming_hosts.json')
+  wvtest.WVFAIL(crash)
+
+
 if __name__ == '__main__':
   wvtest.wvtest_main()