apollo-ios/Tests/ApolloTests/WebSocket/WebSocketTransportTests.swift at 5cd19184cc83245761e5d012be1ef1a586cf19b1 · nabla/apollo-ios · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import XCTest
import Apollo
import ApolloTestSupport
@testable import ApolloWebSocket

class WebSocketTransportTests: XCTestCase {

  private var webSocketTransport: WebSocketTransport!

  override func tearDown() {
    webSocketTransport = nil

    super.tearDown()
  }

  func testUpdateHeaderValues() {
    var request = URLRequest(url: TestURL.mockServer.url)
    request.addValue("OldToken", forHTTPHeaderField: "Authorization")

    self.webSocketTransport = WebSocketTransport(websocket: MockWebSocket(request: request),
                                                 store: ApolloStore())

    self.webSocketTransport.updateHeaderValues(["Authorization": "UpdatedToken"])

    XCTAssertEqual(self.webSocketTransport.websocket.request.allHTTPHeaderFields?["Authorization"], "UpdatedToken")
  }

  func testUpdateConnectingPayload() {
    let request = URLRequest(url: TestURL.mockServer.url)

    self.webSocketTransport = WebSocketTransport(websocket: MockWebSocket(request: request),
                                                 store: ApolloStore(),
                                                 connectingPayload: ["Authorization": "OldToken"])

    let mockWebSocketDelegate = MockWebSocketDelegate()

    let mockWebSocket = self.webSocketTransport.websocket as? MockWebSocket
    self.webSocketTransport.socketConnectionState.mutate { $0 = .connected }
    mockWebSocket?.delegate = mockWebSocketDelegate

    let exp = expectation(description: "Waiting for reconnect")

    mockWebSocketDelegate.didReceiveMessage = { message in
      let json = try? JSONSerializationFormat.deserialize(data: message.data(using: .utf8)!) as? JSONObject
      guard let payload = json?["payload"] as? JSONObject, (json?["type"] as? String) == "connection_init" else {
        return
      }

      XCTAssertEqual(payload["Authorization"] as? String, "UpdatedToken")
      exp.fulfill()
    }

    self.webSocketTransport.updateConnectingPayload(["Authorization": "UpdatedToken"])
    self.webSocketTransport.initServer()

    waitForExpectations(timeout: 3, handler: nil)
  }

  func testCloseConnectionAndInit() {
    let request = URLRequest(url: TestURL.mockServer.url)

    self.webSocketTransport = WebSocketTransport(websocket: MockWebSocket(request: request),
                                                 store: ApolloStore(),
                                                 connectingPayload: ["Authorization": "OldToken"])
    self.webSocketTransport.closeConnection()
    self.webSocketTransport.updateConnectingPayload(["Authorization": "UpdatedToken"])
    self.webSocketTransport.initServer()

    let exp = expectation(description: "Wait")
    let result = XCTWaiter.wait(for: [exp], timeout: 1.0)
    if result == XCTWaiter.Result.timedOut {
    } else {
      XCTFail("Delay interrupted")
    }
  }
}

private final class MockWebSocketDelegate: WebSocketClientDelegate {

  var didReceiveMessage: ((String) -> Void)?

  func websocketDidConnect(socket: WebSocketClient) {}

  func websocketDidDisconnect(socket: WebSocketClient, error: Error?) {}

  func websocketDidReceiveMessage(socket: WebSocketClient, text: String) {
    didReceiveMessage?(text)
  }

  func websocketDidReceiveData(socket: WebSocketClient, data: Data) {}

}