Move WPR-related attributes to UserStorySet.

Make wpr_archive_info property effectively immutable.
Move cloud storage bucket validation logic to page_set_archive_info module.

BUG=439512

Review URL: https://codereview.chromium.org/779383002

Cr-Commit-Position: refs/heads/master@{#307148}

11 files changed, 148 insertions, 98 deletions

diff --git a/tools/perf/measurements/session_restore.py b/tools/perf/measurements/session_restore.py
index 26eef4d..b72bda5 100644
--- a/tools/perf/measurements/session_restore.py
+++ b/tools/perf/measurements/session_restore.py
@@ -56,7 +56,7 @@ class SessionRestore(startup.Startup):
     for page in page_set:
       if page.is_local:
         continue
-      wpr_archive_name = page_set.WprFilePathForPage(page)
+      wpr_archive_name = page_set.WprFilePathForUserStory(page)
       wpr_archive_names_to_page_urls[wpr_archive_name].append(page.url)
 
     # Reject any pageset that contains more than one WPR archive.
diff --git a/tools/telemetry/telemetry/page/__init__.py b/tools/telemetry/telemetry/page/__init__.py
index ab95c21..673a0f1 100644
--- a/tools/telemetry/telemetry/page/__init__.py
+++ b/tools/telemetry/telemetry/page/__init__.py
@@ -203,4 +203,4 @@ class Page(user_story.UserStory):
 
   @property
   def archive_path(self):
-    return self.page_set.WprFilePathForPage(self)
+    return self.page_set.WprFilePathForUserStory(self)
diff --git a/tools/telemetry/telemetry/page/page_set.py b/tools/telemetry/telemetry/page/page_set.py
index 3246ff6..7ad9114 100644
--- a/tools/telemetry/telemetry/page/page_set.py
+++ b/tools/telemetry/telemetry/page/page_set.py
@@ -7,7 +7,6 @@ import inspect
 import os
 
 from telemetry.page import page as page_module
-from telemetry.page import page_set_archive_info
 from telemetry.user_story import user_story_set
 from telemetry.util import cloud_storage
 
@@ -24,9 +23,13 @@ class PageSet(user_story_set.UserStorySet):
   def __init__(self, file_path=None, archive_data_file='', user_agent_type=None,
                make_javascript_deterministic=True, serving_dirs=None,
                bucket=None):
-    super(PageSet, self).__init__()
+    super(PageSet, self).__init__(
+        archive_data_file=archive_data_file, cloud_storage_bucket=bucket)
     # The default value of file_path is location of the file that define this
     # page set instance's class.
+    # TODO(chrishenry): Move this logic to user_story_set. Consider passing
+    # a base_dir directly. Alternatively, kill this and rely on the default
+    # behavior of using the instance's class file location.
     if file_path is None:
       file_path = inspect.getfile(self.__class__)
       # Turn pyc file into py files if we can
@@ -35,17 +38,11 @@ class PageSet(user_story_set.UserStorySet):
 
     self.file_path = file_path
     # These attributes can be set dynamically by the page set.
-    self.archive_data_file = archive_data_file
     self.user_agent_type = user_agent_type
     self.make_javascript_deterministic = make_javascript_deterministic
-    self._wpr_archive_info = None
     # Convert any relative serving_dirs to absolute paths.
     self._serving_dirs = set(os.path.realpath(os.path.join(self.base_dir, d))
                              for d in serving_dirs or [])
-    if self._IsValidPrivacyBucket(bucket):
-      self._bucket = bucket
-    else:
-      raise ValueError("Pageset privacy bucket %s is invalid" % bucket)
 
   @property
   def pages(self):
@@ -66,10 +63,6 @@ class PageSet(user_story_set.UserStorySet):
     self.AddUserStory(page_module.Page(
       page_url, self, self.base_dir))
 
-  @staticmethod
-  def _IsValidPrivacyBucket(bucket_name):
-    return bucket_name in (None, PUBLIC_BUCKET, PARTNER_BUCKET, INTERNAL_BUCKET)
-
   @property
   def base_dir(self):
     if os.path.isfile(self.file_path):
@@ -81,23 +74,6 @@ class PageSet(user_story_set.UserStorySet):
   def serving_dirs(self):
     return self._serving_dirs
 
-  @property
-  def wpr_archive_info(self):  # pylint: disable=E0202
-    """Lazily constructs wpr_archive_info if it's not set and returns it."""
-    if self.archive_data_file and not self._wpr_archive_info:
-      self._wpr_archive_info = (
-          page_set_archive_info.PageSetArchiveInfo.FromFile(
-            os.path.join(self.base_dir, self.archive_data_file), self._bucket))
-    return self._wpr_archive_info
-
-  @property
-  def bucket(self):
-    return self._bucket
-
-  @wpr_archive_info.setter
-  def wpr_archive_info(self, value):  # pylint: disable=E0202
-    self._wpr_archive_info = value
-
   def ContainsOnlyFileURLs(self):
     for page in self.user_stories:
       if not page.is_file:
@@ -127,8 +103,3 @@ class PageSet(user_story_set.UserStorySet):
           raise Exception('Unusable results_file.')
 
     return user_stories
-
-  def WprFilePathForPage(self, page):
-    if not self.wpr_archive_info:
-      return None
-    return self.wpr_archive_info.WprFilePathForPage(page)
diff --git a/tools/telemetry/telemetry/page/page_set_archive_info.py b/tools/telemetry/telemetry/page/page_set_archive_info.py
index 78b4692..eb07998 100644
--- a/tools/telemetry/telemetry/page/page_set_archive_info.py
+++ b/tools/telemetry/telemetry/page/page_set_archive_info.py
@@ -9,11 +9,24 @@ import re
 import shutil
 import tempfile
 
+from telemetry import page as page_module
 from telemetry.util import cloud_storage
 
 
+def AssertValidCloudStorageBucket(bucket):
+  is_valid = bucket in (None,
+                        cloud_storage.PUBLIC_BUCKET,
+                        cloud_storage.PARTNER_BUCKET,
+                        cloud_storage.INTERNAL_BUCKET)
+  if not is_valid:
+    raise ValueError("Cloud storage privacy bucket %s is invalid" % bucket)
+
+
+# TODO(chrishenry): Rename this (and module) to wpr_archive_info.WprArchiveInfo
+# and move to telemetry.user_story or telemetry.wpr or telemetry.core.
 class PageSetArchiveInfo(object):
   def __init__(self, file_path, data, bucket, ignore_archive=False):
+    AssertValidCloudStorageBucket(bucket)
     self._file_path = file_path
     self._base_dir = os.path.dirname(file_path)
     self._bucket = bucket
@@ -63,14 +76,14 @@ class PageSetArchiveInfo(object):
     return cls(file_path, {'archives': {}}, bucket,
                ignore_archive=ignore_archive)
 
-  def WprFilePathForPage(self, page):
+  def WprFilePathForUserStory(self, story):
     if self.temp_target_wpr_file_path:
       return self.temp_target_wpr_file_path
-    wpr_file = self._page_name_to_wpr_file.get(page.display_name, None)
-    if wpr_file is None:
+    wpr_file = self._page_name_to_wpr_file.get(story.display_name, None)
+    if wpr_file is None and isinstance(story, page_module.Page):
       # Some old page sets always use the URL to identify a page rather than the
       # display_name, so try to look for that.
-      wpr_file = self._page_name_to_wpr_file.get(page.url, None)
+      wpr_file = self._page_name_to_wpr_file.get(story.url, None)
     if wpr_file:
       return self._WprFileNameToPath(wpr_file)
     return None
diff --git a/tools/telemetry/telemetry/page/page_set_archive_info_unittest.py b/tools/telemetry/telemetry/page/page_set_archive_info_unittest.py
index bcb4837..4856700 100644
--- a/tools/telemetry/telemetry/page/page_set_archive_info_unittest.py
+++ b/tools/telemetry/telemetry/page/page_set_archive_info_unittest.py
@@ -59,17 +59,17 @@ class TestPageSetArchiveInfo(unittest.TestCase):
       self.assertEquals(cloud_storage.CalculateHash(file_path), f.read())
 
   def testReadingArchiveInfo(self):
-    self.assertIsNotNone(self.archive_info.WprFilePathForPage(page1))
+    self.assertIsNotNone(self.archive_info.WprFilePathForUserStory(page1))
     self.assertEquals(recording1, os.path.basename(
-        self.archive_info.WprFilePathForPage(page1)))
+        self.archive_info.WprFilePathForUserStory(page1)))
 
-    self.assertIsNotNone(self.archive_info.WprFilePathForPage(page2))
+    self.assertIsNotNone(self.archive_info.WprFilePathForUserStory(page2))
     self.assertEquals(recording1, os.path.basename(
-        self.archive_info.WprFilePathForPage(page2)))
+        self.archive_info.WprFilePathForUserStory(page2)))
 
-    self.assertIsNotNone(self.archive_info.WprFilePathForPage(page3))
+    self.assertIsNotNone(self.archive_info.WprFilePathForUserStory(page3))
     self.assertEquals(recording2, os.path.basename(
-        self.archive_info.WprFilePathForPage(page3)))
+        self.archive_info.WprFilePathForUserStory(page3)))
 
   def testArchiveInfoFileGetsUpdated(self):
     """Ensures that the archive info file is updated correctly."""
@@ -106,11 +106,11 @@ class TestPageSetArchiveInfo(unittest.TestCase):
     self.archive_info.AddNewTemporaryRecording(new_temp_recording)
 
     self.assertEquals(new_temp_recording,
-                      self.archive_info.WprFilePathForPage(page1))
+                      self.archive_info.WprFilePathForUserStory(page1))
     self.assertEquals(new_temp_recording,
-                      self.archive_info.WprFilePathForPage(page2))
+                      self.archive_info.WprFilePathForUserStory(page2))
     self.assertEquals(new_temp_recording,
-                      self.archive_info.WprFilePathForPage(page3))
+                      self.archive_info.WprFilePathForUserStory(page3))
 
     self.archive_info.AddRecordedPages([page2])
 
@@ -167,7 +167,7 @@ class TestPageSetArchiveInfo(unittest.TestCase):
     self.archive_info.AddNewTemporaryRecording(new_temp_recording)
 
     self.assertEquals(new_temp_recording,
-                      self.archive_info.WprFilePathForPage(page1))
+                      self.archive_info.WprFilePathForUserStory(page1))
 
     self.archive_info.AddRecordedPages([page1])
 
@@ -183,4 +183,4 @@ class TestPageSetArchiveInfo(unittest.TestCase):
     read_archive_info = page_set_archive_info.PageSetArchiveInfo.FromFile(
         self.page_set_archive_info_file, cloud_storage.PUBLIC_BUCKET)
     self.assertEquals(new_recording,
-                      read_archive_info.WprFilePathForPage(page1))
+                      read_archive_info.WprFilePathForUserStory(page1))
diff --git a/tools/telemetry/telemetry/page/page_set_unittest.py b/tools/telemetry/telemetry/page/page_set_unittest.py
index 6f78121..84ad140 100644
--- a/tools/telemetry/telemetry/page/page_set_unittest.py
+++ b/tools/telemetry/telemetry/page/page_set_unittest.py
@@ -42,25 +42,6 @@ class TestPageSet(unittest.TestCase):
     finally:
       os.rmdir(directory_path)
 
-  def testCloudBucket(self):
-    blank_ps = page_set.PageSet()
-    expected_bucket = None
-    self.assertEqual(blank_ps.bucket, expected_bucket)
-
-    public_ps = page_set.PageSet(bucket=page_set.PUBLIC_BUCKET)
-    expected_bucket = cloud_storage.PUBLIC_BUCKET
-    self.assertEqual(public_ps.bucket, expected_bucket)
-
-    partner_ps = page_set.PageSet(bucket=page_set.PARTNER_BUCKET)
-    expected_bucket = cloud_storage.PARTNER_BUCKET
-    self.assertEqual(partner_ps.bucket, expected_bucket)
-
-    internal_ps = page_set.PageSet(bucket=page_set.INTERNAL_BUCKET)
-    expected_bucket = cloud_storage.INTERNAL_BUCKET
-    self.assertEqual(internal_ps.bucket, expected_bucket)
-
-    self.assertRaises(ValueError, page_set.PageSet, bucket='garbage_bucket')
-
   def testFormingPageSetFromSubPageSet(self):
     page_set_a = page_set.PageSet()
     pages = [
diff --git a/tools/telemetry/telemetry/page/page_test_unittest.py b/tools/telemetry/telemetry/page/page_test_unittest.py
index aac88fd..0654d79 100644
--- a/tools/telemetry/telemetry/page/page_test_unittest.py
+++ b/tools/telemetry/telemetry/page/page_test_unittest.py
@@ -108,7 +108,8 @@ class PageTestUnitTest(page_test_test_case.PageTestTestCase):
       # First record an archive with only www.google.com.
       self._options.browser_options.wpr_mode = wpr_modes.WPR_RECORD
 
-      ps.wpr_archive_info = page_set_archive_info.PageSetArchiveInfo(
+      # pylint: disable=protected-access
+      ps._wpr_archive_info = page_set_archive_info.PageSetArchiveInfo(
           '', '', ps.bucket, json.loads(archive_info_template %
                                         (test_archive, google_url)))
       ps.pages = [page_module.Page(google_url, ps)]
@@ -118,14 +119,16 @@ class PageTestUnitTest(page_test_test_case.PageTestTestCase):
       # Now replay it and verify that google.com is found but foo.com is not.
       self._options.browser_options.wpr_mode = wpr_modes.WPR_REPLAY
 
-      ps.wpr_archive_info = page_set_archive_info.PageSetArchiveInfo(
+      # pylint: disable=protected-access
+      ps._wpr_archive_info = page_set_archive_info.PageSetArchiveInfo(
           '', '', ps.bucket, json.loads(archive_info_template %
                                         (test_archive, foo_url)))
       ps.pages = [page_module.Page(foo_url, ps)]
       all_results = self.RunMeasurement(measurement, ps, options=self._options)
       self.assertEquals(1, len(all_results.failures))
 
-      ps.wpr_archive_info = page_set_archive_info.PageSetArchiveInfo(
+      # pylint: disable=protected-access
+      ps._wpr_archive_info = page_set_archive_info.PageSetArchiveInfo(
           '', '', ps.bucket, json.loads(archive_info_template %
                                         (test_archive, google_url)))
       ps.pages = [page_module.Page(google_url, ps)]
diff --git a/tools/telemetry/telemetry/unittest_util/page_set_smoke_test.py b/tools/telemetry/telemetry/unittest_util/page_set_smoke_test.py
index c7e353c..911a6d8 100644
--- a/tools/telemetry/telemetry/unittest_util/page_set_smoke_test.py
+++ b/tools/telemetry/telemetry/unittest_util/page_set_smoke_test.py
@@ -35,7 +35,7 @@ class PageSetSmokeTest(unittest.TestCase):
     for page in page_set.pages:
       if not page.url.startswith('http'):
         continue
-      self.assertTrue(wpr_archive_info.WprFilePathForPage(page),
+      self.assertTrue(wpr_archive_info.WprFilePathForUserStory(page),
                       msg='No archive found for %s in %s' % (
                           page.url, page_set.archive_data_file))
 
diff --git a/tools/telemetry/telemetry/user_story/user_story_set.py b/tools/telemetry/telemetry/user_story/user_story_set.py
index 8da02a4..a6771fd 100644
--- a/tools/telemetry/telemetry/user_story/user_story_set.py
+++ b/tools/telemetry/telemetry/user_story/user_story_set.py
@@ -2,12 +2,62 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import inspect
+import os
+
 from telemetry import user_story as user_story_module
+from telemetry.page import page_set_archive_info
 
 
 class UserStorySet(object):
-  def __init__(self):
+  """A collection of user story.
+
+  A typical usage of UserStorySet would be to subclass it and then calling
+  AddUserStory for each UserStory..
+  """
+
+  def __init__(self, archive_data_file='', cloud_storage_bucket=None):
+    """Creates a new UserStorySet.
+
+    Args:
+      archive_data_file: The path to Web Page Replay's archive data, relative
+          to self.base_dir.
+      cloud_storage_bucket: The cloud storage bucket used to download
+          Web Page Replay's archive data. Valid values are: None,
+          PUBLIC_BUCKET, PARTNER_BUCKET, or INTERNAL_BUCKET (defined
+          in telemetry.util.cloud_storage).
+    """
     self.user_stories = []
+    self._archive_data_file = archive_data_file
+    self._wpr_archive_info = None
+    page_set_archive_info.AssertValidCloudStorageBucket(cloud_storage_bucket)
+    self._cloud_storage_bucket = cloud_storage_bucket
+    self._base_dir = os.path.dirname(inspect.getfile(self.__class__))
+
+  @property
+  def base_dir(self):
+    """The base directory to resolve archive_data_file.
+
+    This defaults to the directory containing the UserStorySet instance's class.
+    """
+    return self._base_dir
+
+  @property
+  def archive_data_file(self):
+    return self._archive_data_file
+
+  @property
+  def bucket(self):
+    return self._cloud_storage_bucket
+
+  @property
+  def wpr_archive_info(self):
+    """Lazily constructs wpr_archive_info if it's not set and returns it."""
+    if self.archive_data_file and not self._wpr_archive_info:
+      self._wpr_archive_info = (
+          page_set_archive_info.PageSetArchiveInfo.FromFile(
+              os.path.join(self.base_dir, self.archive_data_file), self.bucket))
+    return self._wpr_archive_info
 
   def AddUserStory(self, user_story):
     assert isinstance(user_story, user_story_module.UserStory)
@@ -37,6 +87,11 @@ class UserStorySet(object):
   def ShuffleAndFilterUserStorySet(self, finder_options):
     pass
 
+  def WprFilePathForUserStory(self, story):
+    if not self.wpr_archive_info:
+      return None
+    return self.wpr_archive_info.WprFilePathForUserStory(story)
+
   def __iter__(self):
     return self.user_stories.__iter__()
 
diff --git a/tools/telemetry/telemetry/user_story/user_story_set_unittest.py b/tools/telemetry/telemetry/user_story/user_story_set_unittest.py
new file mode 100644
index 0000000..be7d103
--- /dev/null
+++ b/tools/telemetry/telemetry/user_story/user_story_set_unittest.py
@@ -0,0 +1,50 @@
+# Copyright 2014 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import os
+import unittest
+
+from telemetry.user_story import user_story_set
+from telemetry.util import cloud_storage
+
+
+class UserStorySetFoo(user_story_set.UserStorySet):
+  """ UserStorySetFoo is a user story created for testing purpose. """
+  pass
+
+
+class UserStorySetTest(unittest.TestCase):
+
+  def testUserStoryTestName(self):
+    self.assertEquals('user_story_set_unittest', UserStorySetFoo.Name())
+
+  def testUserStoryTestDescription(self):
+    self.assertEquals(
+        ' UserStorySetFoo is a user story created for testing purpose. ',
+        UserStorySetFoo.Description())
+
+  def testBaseDir(self):
+    uss = UserStorySetFoo()
+    base_dir = uss.base_dir
+    self.assertTrue(os.path.isdir(base_dir))
+    self.assertEqual(base_dir, os.path.dirname(__file__))
+
+  def testCloudBucket(self):
+    blank_uss = user_story_set.UserStorySet()
+    self.assertEqual(blank_uss.bucket, None)
+
+    public_uss = user_story_set.UserStorySet(
+        cloud_storage_bucket=cloud_storage.PUBLIC_BUCKET)
+    self.assertEqual(public_uss.bucket, cloud_storage.PUBLIC_BUCKET)
+
+    partner_uss = user_story_set.UserStorySet(
+        cloud_storage_bucket=cloud_storage.PARTNER_BUCKET)
+    self.assertEqual(partner_uss.bucket, cloud_storage.PARTNER_BUCKET)
+
+    internal_uss = user_story_set.UserStorySet(
+        cloud_storage_bucket=cloud_storage.INTERNAL_BUCKET)
+    self.assertEqual(internal_uss.bucket, cloud_storage.INTERNAL_BUCKET)
+
+    self.assertRaises(ValueError, user_story_set.UserStorySet,
+                      cloud_storage_bucket='garbage_bucket')
diff --git a/tools/telemetry/telemetry/user_story/user_story_test_unittest.py b/tools/telemetry/telemetry/user_story/user_story_test_unittest.py
deleted file mode 100644
index b855d9b..0000000
--- a/tools/telemetry/telemetry/user_story/user_story_test_unittest.py
+++ /dev/null
@@ -1,23 +0,0 @@
-# Copyright 2014 The Chromium Authors. All rights reserved.
-# Use of this source code is governed by a BSD-style license that can be
-# found in the LICENSE file.
-
-import unittest
-
-from telemetry.user_story import user_story_set
-
-
-class UserStorySetFoo(user_story_set.UserStorySet):
-  """ UserStorySetFoo is a user story created for testing purpose. """
-  pass
-
-
-class UserStorySetTest(unittest.TestCase):
-
-  def testUserStoryTestName(self):
-    self.assertEquals('user_story_test_unittest', UserStorySetFoo.Name())
-
-  def testUserStoryTestDescription(self):
-    self.assertEquals(
-        ' UserStorySetFoo is a user story created for testing purpose. ',
-        UserStorySetFoo.Description())