diff --git a/tests/lib/screenshot-testing/support/chai-extras.js b/tests/lib/screenshot-testing/support/chai-extras.js
index 5c24340d1aa7d333e5a025e0e1893cff881eb559..ec2d37296d3365cd3b608e1bc40bd0a865f913bf 100644
--- a/tests/lib/screenshot-testing/support/chai-extras.js
+++ b/tests/lib/screenshot-testing/support/chai-extras.js
@@ -187,7 +187,8 @@ chai.Assertion.addChainableMethod('capture', function () {
 
 // add `contains` assertion
 chai.Assertion.addChainableMethod('contains', function () {
-    var url = this.__flags['object'],
+    var self = this,
+        url = this.__flags['object'],
         elementSelector = arguments[0],
         pageSetupFn = arguments[1],
         done = arguments[2];
@@ -205,21 +206,27 @@ chai.Assertion.addChainableMethod('contains', function () {
     }
 
     pageRenderer.capture(null, function (err) {
+        var obj = self._obj,
+            indent = "     ";
+
         if (err) {
-            var indent = "     ";
             err.stack = err.message + "\n" + indent + getPageLogsString(pageRenderer.pageLogs, indent);
 
             done(err);
             return;
         }
 
-        if (!pageRenderer.contains(elementSelector)) {
-            var error = new AssertionError("Page does not contain element '" + elementSelector + "'.");
-            error.stack = getPageLogsString(pageRenderer.pageLogs, indent);
+        try {
+            self.assert(
+                pageRenderer.contains(elementSelector),
+                "Expected page to contain element '" + elementSelector + "', but could not find it in page.",
+                "Expected page to not contain element '" + elementSelector + "', but found it in page."
+            );
 
-            done(error);
-        } else {
             done();
+        } catch (error) {
+            error.stack = getPageLogsString(pageRenderer.pageLogs, indent);
+            done(error);
         }
     });
 });
\ No newline at end of file
diff --git a/tests/lib/screenshot-testing/support/page-renderer.js b/tests/lib/screenshot-testing/support/page-renderer.js
index aa0fb12f87de441c6d76805424a53e2661657ea5..de3819a4d348fe5efc4c0e8fb6ee0a5895fe0b59 100644
--- a/tests/lib/screenshot-testing/support/page-renderer.js
+++ b/tests/lib/screenshot-testing/support/page-renderer.js
@@ -38,7 +38,7 @@ PageRenderer.prototype.setViewportSize = function (w, h) {
 };
 
 PageRenderer.prototype.getCurrentUrl = function () {
-    return this.webpage.url;
+    return this.webpage ? this.webpage.url : null;
 };
 
 // event queueing functions