2 files changed, 5 insertions, 2 deletions
diff --git a/chrome/content/unittest/run.html b/chrome/content/unittest/run.html
index 1557610..57970d1 100644
--- a/chrome/content/unittest/run.html
+++ b/chrome/content/unittest/run.html
@@ -16,7 +16,7 @@
 
 		//Global variables
 		var subscriptLoader = Components.classes["@mozilla.org/moz/jssubscript-loader;1"].getService(Components.interfaces.mozIJSSubScriptLoader);
-		var redirector = Components.classes["@einaregilsson.com/redirector;1"].getService(Components.interfaces.nsISupports).wrappedJSObject;
+		var redirector = Components.classes["@einaregilsson.com/redirector;1"].getService(Components.interfaces.rdIRedirector);
 		
 		function setupTest(name, testcase) {
 			var table = document.createElement('table');
diff --git a/chrome/content/unittest/testcases.js b/chrome/content/unittest/testcases.js
index 0ce071e..afefff7 100644
--- a/chrome/content/unittest/testcases.js
+++ b/chrome/content/unittest/testcases.js
@@ -95,7 +95,8 @@ var tests = {
 				
 				var ioService = Components.classes["@mozilla.org/network/io-service;1"].getService(Components.interfaces.nsIIOService);  
 				args.contentLocation = ioService.newURI(args.contentLocation, null, null);
-				var result = redirector.shouldLoad(args.contentType, args.contentLocation, args.requestOrigin, args.aContext, args.mimeTypeGuess, args.extra);
+				var contentPolicy = redirector.QueryInterface(nsIContentPolicy);
+				var result = contentPolicy.shouldLoad(args.contentType, args.contentLocation, args.requestOrigin, args.aContext, args.mimeTypeGuess, args.extra);
 				return { passed: result == nsIContentPolicy.ACCEPT, message : "Expected nsIContentPolicy.ACCEPT, actual was " + result };
 			}
 			
@@ -116,6 +117,8 @@ var tests = {
 				try {
 					redirector.enabled = false;
 					return doFunc();
+					redirector.enabled = true;
+
 				} catch(e) {
 					redirector.enabled = true;
 					throw e;