Make sure snap_page combined iframe serialized dom

telemetry/telemetry/internal/snap_page_util.py

Index: telemetry/telemetry/internal/snap_page_util.py
diff --git a/telemetry/telemetry/internal/snap_page_util.py b/telemetry/telemetry/internal/snap_page_util.py
index 57210030776187a0ebe0b2ee0350d5fc817607c3..bef1e129b2263167f4f812b348ca2fc0b4bbf8ac 100644
--- a/telemetry/telemetry/internal/snap_page_util.py
+++ b/telemetry/telemetry/internal/snap_page_util.py
@@ -3,6 +3,7 @@
 # found in the LICENSE file.
 
 import os
+import json
 
 from telemetry.core import util
 from telemetry.internal.browser import browser_finder
@@ -17,28 +18,53 @@ def SnapPage(finder_options, url, interactive, snapshot_file):
   try:
     tab = browser.tabs[0]
     tab.Navigate(url)
-    tab.WaitForDocumentReadyStateToBeComplete()
     if interactive:
       raw_input(
           'Activating interactive mode. Press enter after you finish '
           "interacting with the page to snapshot the page's DOM content.")
-    with open(
-        os.path.join(util.GetTelemetryThirdPartyDir(), 'snap-it',
-                     'HTMLSerializer.js')) as f:
-      snapit_script = f.read()
-    tab.ExecuteJavaScript(snapit_script)
-    tab.ExecuteJavaScript(
-        '''
-        var serializedDomArray;
-        var htmlSerializer = new HTMLSerializer();
-        htmlSerializer.processDocument(document);
-        htmlSerializer.fillHolesAsync(document, function(s) {
-          serializedDomArray = s.html;
-        });
-        ''')
+
     print 'Snapshotting content of %s. This could take a while...' % url
-    tab.WaitForJavaScriptCondition('serializedDomArray !== undefined')
-    serialized_dom = ''.join(tab.EvaluateJavaScript('serializedDomArray'))
-    snapshot_file.write(serialized_dom)
+    tab.WaitForDocumentReadyStateToBeComplete()
+    tab.action_runner.WaitForNetworkQuiescence()
+
+    with open(os.path.join(util.GetTelemetryThirdPartyDir(), 'snap-it',
+                           'HTMLSerializer.js')) as f:
+      snapit_script = f.read()
+
+    with open(os.path.join(util.GetTelemetryThirdPartyDir(), 'snap-it',
+                           'popup.js')) as f:
+      dom_combining_script = f.read()
+
+    serialized_doms = []
+
+    # Serialize the dom in each frame.
+    for context_id in tab.EnableAllContexts():
+      tab.ExecuteJavaScript(snapit_script, context_id=context_id)
+      tab.ExecuteJavaScript(
+          '''
+          var serializedDom;
+          var htmlSerializer = new HTMLSerializer();
+          htmlSerializer.processDocument(document);
+          htmlSerializer.fillHolesAsync(document, function(s) {
+            serializedDom = s.asDict();
+          });
+          ''', context_id=context_id)
+      tab.WaitForJavaScriptCondition(
+          'serializedDom !== undefined', context_id=context_id)
+      serialized_doms.append(tab.EvaluateJavaScript(
+          'serializedDom', context_id=context_id))
+
+    # Sending all the serialized doms back to tab execution context.
+    tab.ExecuteJavaScript('var serializedDoms = [];')
+    for sub_dom in serialized_doms:
+      print len(json.dumps(sub_dom))
+      tab.ExecuteJavaScript('serializedDoms.push({{sub_dom}})', sub_dom=sub_dom)
+
+    # Combine all the doms to one HTML string.
+    tab.EvaluateJavaScript(dom_combining_script)
+    page_snapshot = tab.ExecuteJavaScript(
+        'var domString = outputHTMLString(serialized_doms);')
+
+    snapshot_file.write(page_snapshot)
   finally:
     browser.Close()