1# Copyright 2014 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import inspect
6import os
7
8from telemetry.story import story as story_module
9from telemetry.wpr import archive_info
10
11
12class StorySet(object):
13  """A collection of stories.
14
15  A typical usage of StorySet would be to subclass it and then call
16  AddStory for each Story.
17  """
18
19  def __init__(self, archive_data_file='', cloud_storage_bucket=None,
20               base_dir=None, serving_dirs=None):
21    """Creates a new StorySet.
22
23    Args:
24      archive_data_file: The path to Web Page Replay's archive data, relative
25          to self.base_dir.
26      cloud_storage_bucket: The cloud storage bucket used to download
27          Web Page Replay's archive data. Valid values are: None,
28          story.PUBLIC_BUCKET, story.PARTNER_BUCKET, or story.INTERNAL_BUCKET
29          (defined in telemetry.util.cloud_storage).
30      serving_dirs: A set of paths, relative to self.base_dir, to directories
31          containing hash files for non-wpr archive data stored in cloud
32          storage.
33    """
34    self._stories = []
35    self._story_names_and_grouping_keys = set()
36    self._archive_data_file = archive_data_file
37    self._wpr_archive_info = None
38    archive_info.AssertValidCloudStorageBucket(cloud_storage_bucket)
39    self._cloud_storage_bucket = cloud_storage_bucket
40    if base_dir:
41      if not os.path.isdir(base_dir):
42        raise ValueError('Invalid directory path of base_dir: %s' % base_dir)
43      self._base_dir = base_dir
44    else:
45      self._base_dir = os.path.dirname(inspect.getfile(self.__class__))
46    # Convert any relative serving_dirs to absolute paths.
47    self._serving_dirs = set(os.path.realpath(os.path.join(self.base_dir, d))
48                             for d in serving_dirs or [])
49
50  @property
51  def allow_mixed_story_states(self):
52    """True iff Stories are allowed to have different StoryState classes.
53
54    There are no checks in place for determining if SharedStates are
55    being assigned correctly to all Stories in a given StorySet. The
56    majority of test cases should not need the ability to have multiple
57    SharedStates, which usually implies you should be writing multiple
58    benchmarks instead. We provide errors to avoid accidentally assigning
59    or defaulting to the wrong SharedState.
60    Override at your own risk. Here be dragons.
61    """
62    return False
63
64  @property
65  def file_path(self):
66    return inspect.getfile(self.__class__).replace('.pyc', '.py')
67
68  @property
69  def base_dir(self):
70    """The base directory to resolve archive_data_file.
71
72    This defaults to the directory containing the StorySet instance's class.
73    """
74    return self._base_dir
75
76  @property
77  def serving_dirs(self):
78    all_serving_dirs = self._serving_dirs.copy()
79    for story in self.stories:
80      if story.serving_dir:
81        all_serving_dirs.add(story.serving_dir)
82    return all_serving_dirs
83
84  @property
85  def archive_data_file(self):
86    return self._archive_data_file
87
88  @property
89  def bucket(self):
90    return self._cloud_storage_bucket
91
92  @property
93  def wpr_archive_info(self):
94    """Lazily constructs wpr_archive_info if it's not set and returns it."""
95    if self.archive_data_file and not self._wpr_archive_info:
96      self._wpr_archive_info = archive_info.WprArchiveInfo.FromFile(
97          os.path.join(self.base_dir, self.archive_data_file), self.bucket)
98    return self._wpr_archive_info
99
100  @property
101  def stories(self):
102    return self._stories
103
104  def AddStory(self, story):
105    assert isinstance(story, story_module.Story)
106    assert self._IsUnique(story), ('Tried to add story with duplicate display '
107                                   'name %s. Story display names should be '
108                                   'unique.' % story.display_name)
109    self._stories.append(story)
110    self._story_names_and_grouping_keys.add(
111        story.display_name_and_grouping_key_tuple)
112
113  def _IsUnique(self, story):
114    return (story.display_name_and_grouping_key_tuple not in
115            self._story_names_and_grouping_keys)
116
117  def RemoveStory(self, story):
118    """Removes a Story.
119
120    Allows the stories to be filtered.
121    """
122    self._stories.remove(story)
123    self._story_names_and_grouping_keys.remove(
124        story.display_name_and_grouping_key_tuple)
125
126  @classmethod
127  def Name(cls):
128    """Returns the string name of this StorySet.
129    Note that this should be a classmethod so the benchmark_runner script can
130    match the story class with its name specified in the run command:
131    'Run <User story test name> <User story class name>'
132    """
133    return cls.__module__.split('.')[-1]
134
135  @classmethod
136  def Description(cls):
137    """Return a string explaining in human-understandable terms what this
138    story represents.
139    Note that this should be a classmethod so the benchmark_runner script can
140    display stories' names along with their descriptions in the list command.
141    """
142    if cls.__doc__:
143      return cls.__doc__.splitlines()[0]
144    else:
145      return ''
146
147  def WprFilePathForStory(self, story, target_platform=None):
148    """Convenient function to retrieve WPR archive file path.
149
150    Args:
151      story: The Story to look up.
152
153    Returns:
154      The WPR archive file path for the given Story, if found.
155      Otherwise, None.
156    """
157    if not self.wpr_archive_info:
158      return None
159    return self.wpr_archive_info.WprFilePathForStory(
160        story, target_platform=target_platform)
161
162  def __iter__(self):
163    return self.stories.__iter__()
164
165  def __len__(self):
166    return len(self.stories)
167
168  def __getitem__(self, key):
169    return self.stories[key]
170
171  def __setitem__(self, key, value):
172    self._stories[key] = value
173