1 files changed, 20 insertions, 2 deletions
diff --git a/tests/cpp-unit-tests/browser_mock.cc b/tests/cpp-unit-tests/browser_mock.cc
index 040910b..6b01224 100644
--- a/tests/cpp-unit-tests/browser_mock.cc
+++ b/tests/cpp-unit-tests/browser_mock.cc
@@ -37,7 +37,7 @@
 // Browser mock functions. Add more as needed.
 
 #include <cstring>
-#include <set>
+#include "checked_allocations.h"
 
 #include "UnitTest++.h"
 
@@ -45,7 +45,7 @@
 
 #include "IcedTeaNPPlugin.h"
 
-static std::set<void*> __allocations;
+static AllocationSet __allocations;
 
 // It is expected that these will only run during a unit test
 static void* mock_memalloc(uint32_t size) {
@@ -63,11 +63,29 @@ static void mock_memfree(void* ptr) {
     }
 }
 
+static NPObject* mock_retainobject(NPObject* obj) {
+    obj->referenceCount++;
+    return obj;
+}
+
+static void mock_releaseobject(NPObject* obj) {
+    if (--(obj->referenceCount) == 0) {
+        if (obj->_class->deallocate) {
+            obj->_class->deallocate(obj);
+        } else {
+            free(obj);
+        }
+    }
+}
+
 void browsermock_setup_functions() {
     memset(&browser_functions, 0, sizeof(NPNetscapeFuncs));
 
     browser_functions.memalloc = &mock_memalloc;
     browser_functions.memfree = &mock_memfree;
+
+    browser_functions.retainobject = &mock_retainobject;
+    browser_functions.releaseobject= &mock_releaseobject;
 }
 
 void browsermock_clear_state() {