1 files changed, 72 insertions, 44 deletions

diff --git a/Tools/WebKitTestRunner/WebNotificationProvider.cpp b/Tools/WebKitTestRunner/WebNotificationProvider.cpp
index 55064249f..b009b36a3 100644
--- a/Tools/WebKitTestRunner/WebNotificationProvider.cpp
+++ b/Tools/WebKitTestRunner/WebNotificationProvider.cpp
@@ -26,10 +26,11 @@
 #include "config.h"
 #include "WebNotificationProvider.h"
 
-#include <WebKit2/WKMutableArray.h>
-#include <WebKit2/WKNotification.h>
-#include <WebKit2/WKNumber.h>
-#include <WebKit2/WKSecurityOrigin.h>
+#include <WebKit/WKMutableArray.h>
+#include <WebKit/WKNotification.h>
+#include <WebKit/WKNotificationManager.h>
+#include <WebKit/WKNumber.h>
+#include <WebKit/WKSecurityOriginRef.h>
 #include <wtf/Assertions.h>
 
 namespace WTR {
@@ -65,13 +66,14 @@ WebNotificationProvider::WebNotificationProvider()
 
 WebNotificationProvider::~WebNotificationProvider()
 {
-    WKNotificationManagerSetProvider(m_notificationManager.get(), 0);
+    for (auto& manager : m_ownedNotifications)
+        WKNotificationManagerSetProvider(manager.key.get(), nullptr);
 }
 
 WKNotificationProviderV0 WebNotificationProvider::provider()
 {
     WKNotificationProviderV0 notificationProvider = {
-        { kWKNotificationProviderCurrentVersion, this },
+        { 0, this },
         WTR::showWebNotification,
         WTR::closeWebNotification,
         0, // didDestroyNotification
@@ -83,44 +85,73 @@ WKNotificationProviderV0 WebNotificationProvider::provider()
     return notificationProvider;
 }
 
-void WebNotificationProvider::showWebNotification(WKPageRef, WKNotificationRef notification)
+void WebNotificationProvider::showWebNotification(WKPageRef page, WKNotificationRef notification)
 {
-    if (!m_notificationManager)
-        return;
-
+    auto context = WKPageGetContext(page);
+    auto notificationManager = WKContextGetNotificationManager(context);
     uint64_t id = WKNotificationGetID(notification);
-    ASSERT(!m_shownNotifications.contains(id));
-    m_shownNotifications.add(id);
 
-    WKNotificationManagerProviderDidShowNotification(m_notificationManager.get(), WKNotificationGetID(notification));
+    ASSERT(m_ownedNotifications.contains(notificationManager));
+    auto addResult = m_ownedNotifications.find(notificationManager)->value.add(id);
+    ASSERT_UNUSED(addResult, addResult.isNewEntry);
+    auto addResult2 = m_owningManager.set(id, notificationManager);
+    ASSERT_UNUSED(addResult2, addResult2.isNewEntry);
+    auto addResult3 = m_localToGlobalNotificationIDMap.add(std::make_pair(page, WKNotificationManagerGetLocalIDForTesting(notificationManager, notification)), id);
+    ASSERT_UNUSED(addResult3, addResult3.isNewEntry);
+
+    WKNotificationManagerProviderDidShowNotification(notificationManager, id);
 }
 
-void WebNotificationProvider::closeWebNotification(WKNotificationRef notification)
+static void removeGlobalIDFromIDMap(HashMap<std::pair<WKPageRef, uint64_t>, uint64_t>& map, uint64_t id)
 {
-    if (!m_notificationManager)
-        return;
+    for (auto iter = map.begin(); iter != map.end(); ++iter) {
+        if (iter->value == id) {
+            map.remove(iter);
+            return;
+        }
+    }
+    ASSERT_NOT_REACHED();
+}
 
+void WebNotificationProvider::closeWebNotification(WKNotificationRef notification)
+{
     uint64_t id = WKNotificationGetID(notification);
+    ASSERT(m_owningManager.contains(id));
+    auto notificationManager = m_owningManager.get(id);
+
+    ASSERT(m_ownedNotifications.contains(notificationManager));
+    bool success = m_ownedNotifications.find(notificationManager)->value.remove(id);
+    ASSERT_UNUSED(success, success);
+    m_owningManager.remove(id);
+
+    removeGlobalIDFromIDMap(m_localToGlobalNotificationIDMap, id);
+
     WKRetainPtr<WKUInt64Ref> wkID = WKUInt64Create(id);
     WKRetainPtr<WKMutableArrayRef> array(AdoptWK, WKMutableArrayCreate());
     WKArrayAppendItem(array.get(), wkID.get());
-    m_shownNotifications.remove(id);
-    WKNotificationManagerProviderDidCloseNotifications(m_notificationManager.get(), array.get());
+    WKNotificationManagerProviderDidCloseNotifications(notificationManager, array.get());
 }
 
 void WebNotificationProvider::addNotificationManager(WKNotificationManagerRef manager)
 {
-    // We assume there is only one for testing.
-    ASSERT(!m_notificationManager);
-    m_notificationManager = manager;
+    m_ownedNotifications.add(manager, HashSet<uint64_t>());
 }
 
 void WebNotificationProvider::removeNotificationManager(WKNotificationManagerRef manager)
 {
-    // We assume there is only one for testing.
-    ASSERT(m_notificationManager);
-    ASSERT(m_notificationManager.get() == manager);
-    m_notificationManager = 0;
+    auto iterator = m_ownedNotifications.find(manager);
+    ASSERT(iterator != m_ownedNotifications.end());
+    auto toRemove = iterator->value;
+    WKRetainPtr<WKNotificationManagerRef> guard(manager);
+    m_ownedNotifications.remove(iterator);
+    WKRetainPtr<WKMutableArrayRef> array = adoptWK(WKMutableArrayCreate());
+    for (uint64_t notificationID : toRemove) {
+        bool success = m_owningManager.remove(notificationID);
+        ASSERT_UNUSED(success, success);
+        removeGlobalIDFromIDMap(m_localToGlobalNotificationIDMap, notificationID);
+        WKArrayAppendItem(array.get(), adoptWK(WKUInt64Create(notificationID)).get());
+    }
+    WKNotificationManagerProviderDidCloseNotifications(manager, array.get());
 }
 
 WKDictionaryRef WebNotificationProvider::notificationPermissions()
@@ -129,31 +160,28 @@ WKDictionaryRef WebNotificationProvider::notificationPermissions()
     return WKMutableDictionaryCreate();
 }
 
-void WebNotificationProvider::simulateWebNotificationClick(uint64_t notificationID)
+void WebNotificationProvider::simulateWebNotificationClick(WKPageRef page, uint64_t notificationID)
 {
-    if (!m_notificationManager)
-        return;
-
-    ASSERT(m_shownNotifications.contains(notificationID));
-    WKNotificationManagerProviderDidClickNotification(m_notificationManager.get(), notificationID);
+    ASSERT(m_localToGlobalNotificationIDMap.contains(std::make_pair(page, notificationID)));
+    auto globalID = m_localToGlobalNotificationIDMap.get(std::make_pair(page, notificationID));
+    ASSERT(m_owningManager.contains(globalID));
+    WKNotificationManagerProviderDidClickNotification(m_owningManager.get(globalID), globalID);
 }
 
 void WebNotificationProvider::reset()
 {
-    if (!m_notificationManager) {
-        m_shownNotifications.clear();
-        return;
+    for (auto& notificationPair : m_ownedNotifications) {
+        if (notificationPair.value.isEmpty())
+            continue;
+        WKRetainPtr<WKMutableArrayRef> array = adoptWK(WKMutableArrayCreate());
+        for (uint64_t notificationID : notificationPair.value)
+            WKArrayAppendItem(array.get(), adoptWK(WKUInt64Create(notificationID)).get());
+
+        notificationPair.value.clear();
+        WKNotificationManagerProviderDidCloseNotifications(notificationPair.key.get(), array.get());
     }
-
-    WKRetainPtr<WKMutableArrayRef> array(AdoptWK, WKMutableArrayCreate());
-    HashSet<uint64_t>::const_iterator itEnd = m_shownNotifications.end();
-    for (HashSet<uint64_t>::const_iterator it = m_shownNotifications.begin(); it != itEnd; ++it) {
-        WKRetainPtr<WKUInt64Ref> wkID = WKUInt64Create(*it);
-        WKArrayAppendItem(array.get(), wkID.get());
-    }
-
-    m_shownNotifications.clear();
-    WKNotificationManagerProviderDidCloseNotifications(m_notificationManager.get(), array.get());
+    m_owningManager.clear();
+    m_localToGlobalNotificationIDMap.clear();
 }
 
 } // namespace WTR