Ownership into GenericSets

telemetry/telemetry/internal/story_runner_unittest.py

Index: telemetry/telemetry/internal/story_runner_unittest.py
diff --git a/telemetry/telemetry/internal/story_runner_unittest.py b/telemetry/telemetry/internal/story_runner_unittest.py
index 116b24c6fc901f4747175bdd48cc360e6e2ee21b..f61a362fddf58aa8aa17f2031b9e4b7359a3bb80 100644
--- a/telemetry/telemetry/internal/story_runner_unittest.py
+++ b/telemetry/telemetry/internal/story_runner_unittest.py
@@ -1296,7 +1296,7 @@ class StoryRunnerTest(unittest.TestCase):
     finally:
       shutil.rmtree(temp_path)
 
-  def testRunBenchmark_AddsOwnership_WithoutComponent(self):
+  def testRunBenchmark_AddsOwners_AndEmptyComponent(self):
     @benchmark.Owner(emails=['alice@chromium.org'])
     class FakeBenchmarkWithOwner(FakeBenchmark):
       def __init__(self):
@@ -1318,20 +1318,21 @@ class StoryRunnerTest(unittest.TestCase):
       hs = histogram_set.HistogramSet()
       hs.ImportDicts(data)
 
-      ownership_diagnostics = hs.GetSharedDiagnosticsOfType(
-        histogram_module.Ownership)
+      generic_diagnostics = hs.GetSharedDiagnosticsOfType(
+        histogram_module.GenericSet)
 
-      self.assertGreater(len(ownership_diagnostics), 0)
+      self.assertGreater(len(generic_diagnostics), 0)
 
-      ownership_diagnostic = ownership_diagnostics[0]
+      generic_diagnostics_values = [
+          diagnostic.AsDict()['values'] for diagnostic in generic_diagnostics]
+
+      self.assertIn([], generic_diagnostics_values)
+      self.assertIn(['alice@chromium.org'], generic_diagnostics_values)
 
-      self.assertIsInstance(ownership_diagnostic, histogram_module.Ownership)
-      self.assertIsNone(ownership_diagnostic.component)
-      self.assertItemsEqual(['alice@chromium.org'], ownership_diagnostic.emails)
     finally:
       shutil.rmtree(temp_path)
 
-  def testRunBenchmark_AddsOwnership_WithComponent(self):
+  def testRunBenchmark_AddsComponent(self):
     @benchmark.Owner(emails=['alice@chromium.org', 'bob@chromium.org'],
                      component='fooBar')
     class FakeBenchmarkWithOwner(FakeBenchmark):
@@ -1354,17 +1355,18 @@ class StoryRunnerTest(unittest.TestCase):
       hs = histogram_set.HistogramSet()
       hs.ImportDicts(data)
 
-      ownership_diagnostics = hs.GetSharedDiagnosticsOfType(
-        histogram_module.Ownership)
+      generic_diagnostics = hs.GetSharedDiagnosticsOfType(
+        histogram_module.GenericSet)
+
+      self.assertGreater(len(generic_diagnostics), 0)
 
-      self.assertGreater(len(ownership_diagnostics), 0)
+      generic_diagnostics_values = [
+          diagnostic.AsDict()['values'] for diagnostic in generic_diagnostics]
 
-      ownership_diagnostic = ownership_diagnostics[0]
+      self.assertIn(['fooBar'], generic_diagnostics_values)
+      self.assertIn(['alice@chromium.org', 'bob@chromium.org'],
+                    generic_diagnostics_values)
 
-      self.assertIsInstance(ownership_diagnostic, histogram_module.Ownership)
-      self.assertEqual('fooBar', ownership_diagnostic.component)
-      self.assertItemsEqual(['alice@chromium.org', 'bob@chromium.org'],
-                            ownership_diagnostic.emails)
     finally:
       shutil.rmtree(temp_path)