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._archive_data_file = archive_data_file
36    self._wpr_archive_info = None
37    archive_info.AssertValidCloudStorageBucket(cloud_storage_bucket)
38    self._cloud_storage_bucket = cloud_storage_bucket
39    if base_dir:
40      if not os.path.isdir(base_dir):
41        raise ValueError('Invalid directory path of base_dir: %s' % base_dir)
42      self._base_dir = base_dir
43    else:
44      self._base_dir = os.path.dirname(inspect.getfile(self.__class__))
45    # Convert any relative serving_dirs to absolute paths.
46    self._serving_dirs = set(os.path.realpath(os.path.join(self.base_dir, d))
47                             for d in serving_dirs or [])
48
49  @property
50  def allow_mixed_story_states(self):
51    """True iff Stories are allowed to have different StoryState classes.
52
53    There are no checks in place for determining if SharedStates are
54    being assigned correctly to all Stories in a given StorySet. The
55    majority of test cases should not need the ability to have multiple
56    SharedStates, which usually implies you should be writing multiple
57    benchmarks instead. We provide errors to avoid accidentally assigning
58    or defaulting to the wrong SharedState.
59    Override at your own risk. Here be dragons.
60    """
61    return False
62
63  @property
64  def file_path(self):
65    return inspect.getfile(self.__class__).replace('.pyc', '.py')
66
67  @property
68  def base_dir(self):
69    """The base directory to resolve archive_data_file.
70
71    This defaults to the directory containing the StorySet instance's class.
72    """
73    return self._base_dir
74
75  @property
76  def serving_dirs(self):
77    all_serving_dirs = self._serving_dirs.copy()
78    for story in self.stories:
79      if story.serving_dir:
80        all_serving_dirs.add(story.serving_dir)
81    return all_serving_dirs
82
83  @property
84  def archive_data_file(self):
85    return self._archive_data_file
86
87  @property
88  def bucket(self):
89    return self._cloud_storage_bucket
90
91  @property
92  def wpr_archive_info(self):
93    """Lazily constructs wpr_archive_info if it's not set and returns it."""
94    if self.archive_data_file and not self._wpr_archive_info:
95      self._wpr_archive_info = archive_info.WprArchiveInfo.FromFile(
96          os.path.join(self.base_dir, self.archive_data_file), self.bucket)
97    return self._wpr_archive_info
98
99  def AddStory(self, story):
100    assert isinstance(story, story_module.Story)
101    self.stories.append(story)
102
103  def RemoveStory(self, story):
104    """Removes a Story.
105
106    Allows the stories to be filtered.
107    """
108    self.stories.remove(story)
109
110  @classmethod
111  def Name(cls):
112    """Returns the string name of this StorySet.
113    Note that this should be a classmethod so the benchmark_runner script can
114    match the story class with its name specified in the run command:
115    'Run <User story test name> <User story class name>'
116    """
117    return cls.__module__.split('.')[-1]
118
119  @classmethod
120  def Description(cls):
121    """Return a string explaining in human-understandable terms what this
122    story represents.
123    Note that this should be a classmethod so the benchmark_runner script can
124    display stories' names along with their descriptions in the list command.
125    """
126    if cls.__doc__:
127      return cls.__doc__.splitlines()[0]
128    else:
129      return ''
130
131  def WprFilePathForStory(self, story):
132    """Convenient function to retrieve WPR archive file path.
133
134    Args:
135      story: The Story to look up.
136
137    Returns:
138      The WPR archive file path for the given Story, if found.
139      Otherwise, None.
140    """
141    if not self.wpr_archive_info:
142      return None
143    return self.wpr_archive_info.WprFilePathForStory(story)
144
145  def __iter__(self):
146    return self.stories.__iter__()
147
148  def __len__(self):
149    return len(self.stories)
150
151  def __getitem__(self, key):
152    return self.stories[key]
153
154  def __setitem__(self, key, value):
155    self.stories[key] = value
156