Move WPR-related attributes to UserStorySet.

tools/telemetry/telemetry/page/page_set_archive_info.py

 # Copyright 2013 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 json
 import logging
 import os
 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
 
     # Ensure directory exists.
     if not os.path.exists(self._base_dir):
       os.makedirs(self._base_dir)
 
     # Download all .wpr files.
     if not ignore_archive:
@@ skipping 29 matching lines @@
 
   @classmethod
   def FromFile(cls, file_path, bucket, ignore_archive=False):
     if os.path.exists(file_path):
       with open(file_path, 'r') as f:
         data = json.load(f)
         return cls(file_path, data, bucket, ignore_archive=ignore_archive)
     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
 
   def AddNewTemporaryRecording(self, temp_wpr_file_path=None):
     if temp_wpr_file_path is None:
       temp_wpr_file_handle, temp_wpr_file_path = tempfile.mkstemp()
       os.close(temp_wpr_file_handle)
     self.temp_target_wpr_file_path = temp_wpr_file_path
 
@@ skipping 90 matching lines @@
 
   def _SetWprFileForPage(self, page_name, wpr_file):
     """For modifying the metadata when we're going to record a new archive."""
     old_wpr_file = self._page_name_to_wpr_file.get(page_name, None)
     if old_wpr_file:
       self._wpr_file_to_page_names[old_wpr_file].remove(page_name)
     self._page_name_to_wpr_file[page_name] = wpr_file
     if wpr_file not in self._wpr_file_to_page_names:
       self._wpr_file_to_page_names[wpr_file] = []
     self._wpr_file_to_page_names[wpr_file].append(page_name)