From f458dc81cc9c65ab4340dd4c5e55fabe1e11dd16 Mon Sep 17 00:00:00 2001
From: Squidonomics <eigthmonthstothemoon@gmail.com>
Date: Fri, 23 Feb 2024 18:02:09 -0500
Subject: [PATCH] fix: Address the pr comments

---
 Docs/SupportedAPIs.md                         |  2 +-
 Sources/WebDriver/Requests.swift              | 19 ++++++++++++-------
 Sources/WebDriver/Session.swift               | 12 ++++++++++++
 .../UnitTests/APIToRequestMappingTests.swift  | 15 ++++++++-------
 4 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/Docs/SupportedAPIs.md b/Docs/SupportedAPIs.md
index b7491ed..ba8ca67 100644
--- a/Docs/SupportedAPIs.md
+++ b/Docs/SupportedAPIs.md
@@ -70,4 +70,4 @@ Contributions to expand support to unimplemented functionality are always welcom
 | GET    | `/session/:sessionId/window/:windowHandle/position` | Supported    | Not implemented     |
 | POST   | `/session/:sessionId/window/:windowHandle/maximize` | Supported    | Not implemented     |
 | GET    | `/session/:sessionId/window_handle`                 | Supported    | `Session.windowHandle()`|
-| GET    | `/session/:sessionId/window_handles`                | Supported    | `Session.windowHandle()`|
\ No newline at end of file
+| GET    | `/session/:sessionId/window_handles`                | Supported    | `Session.windowHandles()`|
\ No newline at end of file
diff --git a/Sources/WebDriver/Requests.swift b/Sources/WebDriver/Requests.swift
index 3aafff4..66bf035 100644
--- a/Sources/WebDriver/Requests.swift
+++ b/Sources/WebDriver/Requests.swift
@@ -564,17 +564,22 @@ public enum Requests {
         public typealias Response = WebDriverStatus
     }
 
+    // https://www.selenium.dev/documentation/legacy/json_wire_protocol/#sessionsessionidwindow_handle
     public struct SessionWindowHandle: Request {
         public var session: String 
-        public var many: Bool
 
-        public var pathComponents: [String] { ["session", session, many ? "window_handles" : "window_handle"] }
+        public var pathComponents: [String] { ["session", session, "window_handle"] }
         public var method: HTTPMethod { .get }
 
-        public typealias Response = ResponseWithValue<ResponseValue>
-        public struct ResponseValue: Codable {
-            public var windowHandle: String
-            public var windowHandles: [String]
-        }
+        public typealias Response = ResponseWithValue<String>
+    }
+
+    public struct SessionWindowHandles: Request {
+        public var session: String 
+
+        public var pathComponents: [String] { ["session", session, "window_handles"] }
+        public var method: HTTPMethod { .get }
+
+        public typealias Response = ResponseWithValue<Array<String>>
     }
 }
diff --git a/Sources/WebDriver/Session.swift b/Sources/WebDriver/Session.swift
index b59f4ab..29675f0 100644
--- a/Sources/WebDriver/Session.swift
+++ b/Sources/WebDriver/Session.swift
@@ -326,6 +326,18 @@ public class Session {
         let response = try webDriver.send(Requests.SessionSource(session: id))
         return response.value.source
     }
+ 
+    /// - Returns: Current window handle
+    public func windowHandle() throws -> String {
+        let response = try webDriver.send(Requests.SessionWindowHandle(session: id))
+        return response.value
+    }
+
+    /// - Returns: Array of window handles
+    public func windowHandles() throws ->  [String] {
+        let response = try webDriver.send(Requests.SessionWindowHandles(session: id))
+        return response.value
+    }
 
     /// Deletes the current session.
     public func delete() throws {
diff --git a/Tests/UnitTests/APIToRequestMappingTests.swift b/Tests/UnitTests/APIToRequestMappingTests.swift
index fa8104f..cfe5f11 100644
--- a/Tests/UnitTests/APIToRequestMappingTests.swift
+++ b/Tests/UnitTests/APIToRequestMappingTests.swift
@@ -219,19 +219,20 @@ class APIToRequestMappingTests: XCTestCase {
         let session = Session(webDriver: mockWebDriver, existingId: "mySession")
 
         mockWebDriver.expect(path: "session/mySession/window_handle", method: .get, type: Requests.SessionWindowHandle.self) {
-            ResponseWithValue(.init(windowHandle: "myWindow", windowHandles: []))
+            ResponseWithValue(.init("myWindow"))
         }
-        XCTAssert(try session.windowHandle(many: false) == (windowHandle: "myWindow", windowHandles: []))
+        XCTAssert(try session.windowHandle() == "myWindow")
     }
 
-        func testWindowHandles() throws {
+    func testWindowHandles() throws {
         let mockWebDriver: MockWebDriver = MockWebDriver()
         let session = Session(webDriver: mockWebDriver, existingId: "mySession")
-
-        mockWebDriver.expect(path: "session/mySession/window_handles", method: .get, type: Requests.SessionWindowHandle.self) {
-            ResponseWithValue(.init(windowHandle: "", windowHandles: ["myWindow", "myWindow"]))
+        
+        mockWebDriver.expect(path: "session/mySession/window_handles", method: .get, type: Requests.SessionWindowHandles.self) {
+            ResponseWithValue(.init(["myWindow", "myWindow"]))
         }
-        XCTAssert(try session.windowHandle(many: true) == (windowHandle: "", windowHandles: ["myWindow", "myWindow"]))
+        XCTAssert(try session.windowHandles() == ["myWindow", "myWindow"])
+    }
 
     func testSessionSource() throws {
         let mockWebDriver: MockWebDriver = MockWebDriver()