add unit tests for RelayPool

damus.xcodeproj/project.pbxproj

@@ -202,6 +202,7 @@
 		4CF0ABF029857E9200D66079 /* Bech32Object.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4CF0ABEF29857E9200D66079 /* Bech32Object.swift */; };
 		4CF0ABF62985CD5500D66079 /* UserSearch.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4CF0ABF52985CD5500D66079 /* UserSearch.swift */; };
 		4FE60CDD295E1C5E00105A1F /* Wallet.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4FE60CDC295E1C5E00105A1F /* Wallet.swift */; };
+		5023E76329AA3627007D3D50 /* RelayPoolTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5023E76229AA3627007D3D50 /* RelayPoolTests.swift */; };
 		50A50A8D29A09E1C00C01BE7 /* RequestTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 50A50A8C29A09E1C00C01BE7 /* RequestTests.swift */; };
 		5C513FBA297F72980072348F /* CustomPicker.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5C513FB9297F72980072348F /* CustomPicker.swift */; };
 		5C513FCC2984ACA60072348F /* QRCodeView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5C513FCB2984ACA60072348F /* QRCodeView.swift */; };
@@ -522,6 +523,7 @@
 		4CF0ABEF29857E9200D66079 /* Bech32Object.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Bech32Object.swift; sourceTree = "<group>"; };
 		4CF0ABF52985CD5500D66079 /* UserSearch.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UserSearch.swift; sourceTree = "<group>"; };
 		4FE60CDC295E1C5E00105A1F /* Wallet.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Wallet.swift; sourceTree = "<group>"; };
+		5023E76229AA3627007D3D50 /* RelayPoolTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = RelayPoolTests.swift; sourceTree = "<group>"; };
 		50A50A8C29A09E1C00C01BE7 /* RequestTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = RequestTests.swift; sourceTree = "<group>"; };
 		5C513FB9297F72980072348F /* CustomPicker.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CustomPicker.swift; sourceTree = "<group>"; };
 		5C513FCB2984ACA60072348F /* QRCodeView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = QRCodeView.swift; sourceTree = "<group>"; };
@@ -959,6 +961,7 @@
 		4CE6DEF627F7A08200C66700 /* damusTests */ = {
 			isa = PBXGroup;
 			children = (
+				5023E76229AA3627007D3D50 /* RelayPoolTests.swift */,
 				50A50A8C29A09E1C00C01BE7 /* RequestTests.swift */,
 				DD597CBC2963D85A00C64D32 /* MarkdownTests.swift */,
 				4C90BD1B283AC38E008EE7EF /* Bech32Tests.swift */,
@@ -1419,6 +1422,7 @@
 				3ACBCB78295FE5C70037388A /* TimeAgoTests.swift in Sources */,
 				DD597CBD2963D85A00C64D32 /* MarkdownTests.swift in Sources */,
 				4C3EA67B28FF7B3900C48A62 /* InvoiceTests.swift in Sources */,
+				5023E76329AA3627007D3D50 /* RelayPoolTests.swift in Sources */,
 				4C363A9E2828A822006E126D /* ReplyTests.swift in Sources */,
 				4CB883AA297612FF00DC99E7 /* ZapTests.swift in Sources */,
 				4CB8839A297322D200DC99E7 /* DMTests.swift in Sources */,

damus/Nostr/RelayPool.swift

@@ -18,7 +18,7 @@ struct QueuedRequest {
 }
 
 final class RelayPool {
-    private enum Constants {
+    enum Constants {
         /// Used for an exponential backoff algorithm when retrying stale connections
         /// Each retry attempt will be delayed by raising this base delay to an exponent
         /// equal to the number of previous retries.
@@ -28,10 +28,10 @@ final class RelayPool {
     }
     
     private(set) var relays: [Relay] = []
-    private var handlers: [RelayHandler] = []
+    private(set) var handlers: [RelayHandler] = []
     private var request_queue: [QueuedRequest] = []
-    private var seen: Set<String> = Set()
-    private var counts: [String: UInt64] = [:]
+    private(set) var seen: Set<String> = Set()
+    private(set) var counts: [String: UInt64] = [:]
     private var retry_attempts_per_relay: [URL: Int] = [:]
 
     var descriptors: [RelayDescriptor] {
@@ -42,8 +42,11 @@ final class RelayPool {
         relays.reduce(0) { n, r in n + (r.connection.state == .connecting ? 1 : 0) }
     }
 
-    private func remove_handler(sub_id: String) {
-        self.handlers = handlers.filter { $0.sub_id != sub_id }
+    func remove_handler(sub_id: String) {
+        guard let index = handlers.firstIndex(where: { $0.sub_id == sub_id }) else {
+            return
+        }
+        handlers.remove(at: index)
         print("removing \(sub_id) handler, current: \(handlers.count)")
     }
 
@@ -64,7 +67,8 @@ final class RelayPool {
         }
     }
     
-    func add_relay(_ url: URL, info: RelayInfo) throws {
+    @discardableResult
+    func add_relay(_ url: URL, info: RelayInfo) throws -> Relay {
         let relay_id = get_relay_id(url)
         if get_relay(relay_id) != nil {
             throw RelayError.RelayAlreadyExists
@@ -74,7 +78,8 @@ final class RelayPool {
         }
         let descriptor = RelayDescriptor(url: url, info: info)
         let relay = Relay(descriptor: descriptor, connection: conn)
-        self.relays.append(relay)
+        relays.append(relay)
+        return relay
     }
     
     /// This is used to retry dead connections
@@ -152,13 +157,13 @@ final class RelayPool {
         send(.subscribe(.init(filters: filters, sub_id: sub_id)), to: to)
     }
     
-    private func count_queued(relay: String) -> Int {
+    func count_queued(relay: String) -> Int {
         request_queue.filter({ $0.relay == relay }).count
     }
     
     func queue_req(r: NostrRequest, relay: String) {
         let count = count_queued(relay: relay)
-        guard count <= Constants.max_queued_requests else {
+        guard count < Constants.max_queued_requests else {
             print("can't queue, too many queued events for \(relay)")
             return
         }
@@ -210,7 +215,7 @@ final class RelayPool {
         }
     }
     
-    private func record_seen(relay_id: String, event: NostrConnectionEvent) {
+    func record_seen(relay_id: String, event: NostrConnectionEvent) {
         if case .nostr_event(let ev) = event {
             if case .event(_, let nev) = ev {
                 let k = relay_id + nev.id
@@ -243,5 +248,5 @@ final class RelayPool {
 
 func add_rw_relay(_ pool: RelayPool, _ url: String) {
     let url_ = URL(string: url)!
-    try? pool.add_relay(url_, info: RelayInfo.rw)
+    let _ = try? pool.add_relay(url_, info: RelayInfo.rw)
 }

damusTests/RelayPoolTests.swift

@@ -0,0 +1,170 @@
+//
+//  RelayPoolTests.swift
+//  damusTests
+//
+//  Created by Bryan Montz on 2/25/23.
+//
+
+import XCTest
+@testable import damus
+
+final class RelayPoolTests: XCTestCase {
+    
+    private let fakeRelayURL = URL(string: "wss://some.relay.com")!
+    
+    private func setUpPool() throws -> RelayPool {
+        let pool = RelayPool()
+        XCTAssertTrue(pool.relays.isEmpty)
+        
+        try pool.add_relay(fakeRelayURL, info: RelayInfo.rw)
+        return pool
+    }
+    
+    // MARK: - Relay Add/Remove
+    
+    func testAddRelay() throws {
+        let pool = try setUpPool()
+        
+        XCTAssertEqual(pool.relays.count, 1)
+    }
+    
+    func testRejectDuplicateRelay() throws {
+        let pool = try setUpPool()
+        
+        XCTAssertThrowsError(try pool.add_relay(fakeRelayURL, info: RelayInfo.rw)) { error in
+            XCTAssertEqual(error as? RelayError, RelayError.RelayAlreadyExists)
+        }
+    }
+    
+    func testRemoveRelay() throws {
+        let pool = try setUpPool()
+        
+        XCTAssertEqual(pool.relays.count, 1)
+        
+        pool.remove_relay(fakeRelayURL.absoluteString)
+        
+        XCTAssertTrue(pool.relays.isEmpty)
+    }
+    
+    func testMarkRelayBroken() throws {
+        let pool = try setUpPool()
+        
+        let relay = try XCTUnwrap(pool.relays.first(where: { $0.id == fakeRelayURL.absoluteString }))
+        XCTAssertFalse(relay.is_broken)
+        
+        pool.mark_broken(fakeRelayURL.absoluteString)
+        XCTAssertTrue(relay.is_broken)
+    }
+    
+    func testGetRelay() throws {
+        let pool = try setUpPool()
+        XCTAssertNotNil(pool.get_relay(fakeRelayURL.absoluteString))
+    }
+    
+    func testGetRelays() throws {
+        let pool = try setUpPool()
+        
+        try pool.add_relay(URL(string: "wss://second.relay.com")!, info: RelayInfo.rw)
+        
+        let allRelays = pool.get_relays([fakeRelayURL.absoluteString, "wss://second.relay.com"])
+        XCTAssertEqual(allRelays.count, 2)
+        
+        let relays = pool.get_relays(["wss://second.relay.com"])
+        XCTAssertEqual(relays.count, 1)
+    }
+    
+    // MARK: - Handler Add/Remove
+    
+    private func setUpPoolWithHandler(sub_id: String) -> RelayPool {
+        let pool = RelayPool()
+        XCTAssertTrue(pool.handlers.isEmpty)
+        
+        pool.register_handler(sub_id: sub_id) { _, _ in }
+        return pool
+    }
+    
+    func testAddHandler() {
+        let sub_id = "123"
+        let pool = setUpPoolWithHandler(sub_id: sub_id)
+        
+        XCTAssertEqual(pool.handlers.count, 1)
+    }
+    
+    func testRejectDuplicateHandler() {
+        let sub_id = "123"
+        let pool = setUpPoolWithHandler(sub_id: sub_id)
+        XCTAssertEqual(pool.handlers.count, 1)
+        
+        pool.register_handler(sub_id: sub_id) { _, _ in }
+        
+        XCTAssertEqual(pool.handlers.count, 1)
+    }
+    
+    func testRemoveHandler() {
+        let sub_id = "123"
+        let pool = setUpPoolWithHandler(sub_id: sub_id)
+        XCTAssertEqual(pool.handlers.count, 1)
+        pool.remove_handler(sub_id: sub_id)
+        XCTAssertTrue(pool.handlers.isEmpty)
+    }
+    
+    func testRecordLastPong() throws {
+        let pool = try setUpPool()
+        let relayId = fakeRelayURL.absoluteString
+        let relay = try XCTUnwrap(pool.get_relay(relayId))
+        XCTAssertEqual(relay.last_pong, 0)
+        
+        let pongEvent = NostrConnectionEvent.ws_event(.pong(nil))
+        pool.record_last_pong(relay_id: relayId, event: pongEvent)
+        XCTAssertNotEqual(relay.last_pong, 0)
+    }
+    
+    func testSeenAndCounts() throws {
+        let pool = try setUpPool()
+        
+        XCTAssertTrue(pool.seen.isEmpty)
+        XCTAssertTrue(pool.counts.isEmpty)
+        
+        let event = NostrEvent(id: "123", content: "", pubkey: "")
+        let connectionEvent = NostrConnectionEvent.nostr_event(NostrResponse.event("", event))
+        let relay_id = fakeRelayURL.absoluteString
+        pool.record_seen(relay_id: relay_id, event: connectionEvent)
+        
+        XCTAssertTrue(pool.seen.contains("wss://some.relay.com123"))
+        
+        XCTAssertEqual(pool.counts[relay_id], 1)
+        
+        pool.record_seen(relay_id: relay_id, event: connectionEvent)
+        // don't count the same event twice
+        XCTAssertEqual(pool.counts[relay_id], 1)
+    }
+    
+    func testAddQueuedRequest() throws {
+        let pool = try setUpPool()
+        
+        XCTAssertEqual(pool.count_queued(relay: fakeRelayURL.absoluteString), 0)
+        
+        let req = NostrRequest.unsubscribe("")
+        pool.queue_req(r: req, relay: fakeRelayURL.absoluteString)
+        
+        XCTAssertEqual(pool.count_queued(relay: fakeRelayURL.absoluteString), 1)
+    }
+    
+    func testRejectTooManyQueuedRequests() throws {
+        let pool = try setUpPool()
+        
+        let maxRequests = RelayPool.Constants.max_queued_requests
+        for _ in 0..<maxRequests {
+            let req = NostrRequest.unsubscribe("")
+            pool.queue_req(r: req, relay: fakeRelayURL.absoluteString)
+        }
+        
+        XCTAssertEqual(pool.count_queued(relay: fakeRelayURL.absoluteString), maxRequests)
+        
+        // try to add one beyond the maximum
+        let req = NostrRequest.unsubscribe("")
+        pool.queue_req(r: req, relay: fakeRelayURL.absoluteString)
+        
+        XCTAssertEqual(pool.count_queued(relay: fakeRelayURL.absoluteString), maxRequests)
+    }
+}