diff --git a/Sources/AXorcist/Core/AXorcist.swift b/Sources/AXorcist/Core/AXorcist.swift index c80dcd4..61464ba 100644 --- a/Sources/AXorcist/Core/AXorcist.swift +++ b/Sources/AXorcist/Core/AXorcist.swift @@ -121,6 +121,7 @@ public class AXorcist { // Collect all elements recursively var collectedElements: [AXElementData] = [] + var visitedElements: Set = [] let attributesToFetch = command.attributesToReturn ?? AXMiscConstants.defaultAttributesToFetch let collectionContext = ElementCollectionContext( maxDepth: command.maxDepth, @@ -131,6 +132,7 @@ public class AXorcist { element: rootElement, currentDepth: 0, context: collectionContext, + visited: &visitedElements, collectedElements: &collectedElements) self.logger.log(AXLogEntry( @@ -204,11 +206,16 @@ public class AXorcist { element: Element, currentDepth: Int, context: ElementCollectionContext, + visited: inout Set, collectedElements: inout [AXElementData]) { // Check depth limit guard currentDepth <= context.maxDepth else { return } + // Prevent infinite loops caused by cyclic AX hierarchies (Window → App → Window …) + let hashValue = CFHash(element.underlyingElement) + guard visited.insert(hashValue).inserted else { return } + // Apply filter criteria if provided if let criteria = context.filterCriteria { guard elementMatchesCriteria(element, criteria: criteria) else { return } @@ -228,6 +235,7 @@ public class AXorcist { element: child, currentDepth: currentDepth + 1, context: context, + visited: &visited, collectedElements: &collectedElements) } } diff --git a/Tests/AXorcistTests/ApplicationQueryTests.swift b/Tests/AXorcistTests/ApplicationQueryTests.swift index 880c3ed..665ed7b 100644 --- a/Tests/AXorcistTests/ApplicationQueryTests.swift +++ b/Tests/AXorcistTests/ApplicationQueryTests.swift @@ -38,8 +38,10 @@ struct ApplicationQueryTests { let command = CommandEnvelope( commandId: "test-get-all-apps", command: .collectAll, + attributes: ["AXRole", "AXTitle", "AXIdentifier"], debugLogging: true, locator: Locator(criteria: [Criterion(attribute: "AXRole", value: "AXApplication")]), + maxDepth: 3, outputFormat: .verbose) let encoder = JSONEncoder() @@ -60,13 +62,24 @@ struct ApplicationQueryTests { throw TestError.generic("No output") } - let response = try JSONDecoder().decode(QueryResponse.self, from: responseData) - - #expect(response.success) - #expect(response.data != nil, "Should have data") - - if let data = response.data { - #expect(data.attributes != nil, "Should have attributes") + if let response = try? JSONDecoder().decode(QueryResponse.self, from: responseData) { + #expect(response.success) + if let data = response.data { + #expect(data.attributes != nil, "Should have attributes") + } else { + Issue.record("CollectAll query response had no data payload") + } + } else if + let jsonObject = try? JSONSerialization.jsonObject(with: responseData) as? [String: Any], + let data = jsonObject["data"] as? [String: Any], + let count = data["count"] as? Int, + let elements = data["elements"] as? [[String: Any]] + { + #expect(count > 0, "CollectAll response should report at least one element") + #expect(!elements.isEmpty, "CollectAll response should include element payloads") + } else { + let fallback = String(data: responseData, encoding: .utf8) ?? "" + Issue.record("Unexpected response payload for collectAll: \(fallback)") } } @@ -149,13 +162,21 @@ struct ApplicationQueryTests { throw TestError.generic("No output") } - let response = try JSONDecoder().decode(SimpleSuccessResponse.self, from: responseData) - - if response.success { - let message = response.message + if let response = try? JSONDecoder().decode(SimpleSuccessResponse.self, from: responseData) { + if response.success { + let message = response.message + #expect( + message.contains("No") || message.contains("not found") || message.isEmpty, + "Message should indicate no elements found or be empty") + } + } else if let jsonObject = try? JSONSerialization.jsonObject(with: responseData) as? [String: Any] { + let message = (jsonObject["message"] as? String) ?? (jsonObject["error"] as? String) ?? "" #expect( - message.contains("No") || message.contains("not found") || message.isEmpty, + message.contains("No") || message.contains("not found") || message.contains("error") || message.isEmpty, "Message should indicate no elements found or be empty") + } else { + let rawOutput = String(data: responseData, encoding: .utf8) ?? "" + Issue.record("Unexpected response for nonexistent app: \(rawOutput)") } } } diff --git a/Tests/AXorcistTests/CommonTestHelpers.swift b/Tests/AXorcistTests/CommonTestHelpers.swift index 2493426..5a301c2 100644 --- a/Tests/AXorcistTests/CommonTestHelpers.swift +++ b/Tests/AXorcistTests/CommonTestHelpers.swift @@ -183,11 +183,14 @@ func runAXORCCommand(arguments: [String]) throws -> CommandResult { process.standardOutput = outputPipe process.standardError = errorPipe + let readStdout = startStreaming(pipe: outputPipe) + let readStderr = startStreaming(pipe: errorPipe) + try process.run() process.waitUntilExit() - let outputData = outputPipe.fileHandleForReading.readDataToEndOfFile() - let errorData = errorPipe.fileHandleForReading.readDataToEndOfFile() + let outputData = readStdout() + let errorData = readStderr() let output = String(data: outputData, encoding: String.Encoding.utf8)? .trimmingCharacters(in: .whitespacesAndNewlines) @@ -235,6 +238,9 @@ func runAXORCCommandWithStdin(inputJSON: String, arguments: [String]) throws -> process.standardError = errorPipe process.standardInput = inputPipe + let readStdout = startStreaming(pipe: outputPipe) + let readStderr = startStreaming(pipe: errorPipe) + try process.run() if let inputData = inputJSON.data(using: String.Encoding.utf8) { @@ -247,8 +253,8 @@ func runAXORCCommandWithStdin(inputJSON: String, arguments: [String]) throws -> process.waitUntilExit() - let outputData = outputPipe.fileHandleForReading.readDataToEndOfFile() - let errorData = errorPipe.fileHandleForReading.readDataToEndOfFile() + let outputData = readStdout() + let errorData = readStderr() let output = String(data: outputData, encoding: String.Encoding.utf8)? .trimmingCharacters(in: .whitespacesAndNewlines) @@ -280,6 +286,7 @@ struct CommandEnvelope: Codable { locator: Locator? = nil, pathHint: [String]? = nil, maxElements: Int? = nil, + maxDepth: Int? = nil, outputFormat: OutputFormat? = nil, actionName: String? = nil, actionValue: AttributeValue? = nil, @@ -294,6 +301,7 @@ struct CommandEnvelope: Codable { self.locator = locator self.pathHint = pathHint self.maxElements = maxElements + self.maxDepth = maxDepth self.outputFormat = outputFormat self.actionName = actionName self.actionValue = actionValue @@ -312,6 +320,7 @@ struct CommandEnvelope: Codable { case locator case pathHint = "path_hint" case maxElements = "max_elements" + case maxDepth = "max_depth" case outputFormat = "output_format" case actionName = "action_name" case actionValue = "action_value" @@ -327,6 +336,7 @@ struct CommandEnvelope: Codable { let locator: Locator? let pathHint: [String]? let maxElements: Int? + let maxDepth: Int? let outputFormat: OutputFormat? let actionName: String? let actionValue: AttributeValue? diff --git a/Tests/AXorcistTests/Support/PipeStreaming.swift b/Tests/AXorcistTests/Support/PipeStreaming.swift new file mode 100644 index 0000000..55cbbae --- /dev/null +++ b/Tests/AXorcistTests/Support/PipeStreaming.swift @@ -0,0 +1,26 @@ +import Foundation + +@discardableResult +func startStreaming(pipe: Pipe) -> () -> Data { + let queue = DispatchQueue(label: "axorcist.tests.pipe-stream.\(UUID().uuidString)", qos: .userInitiated) + var collected = Data() + let group = DispatchGroup() + + group.enter() + pipe.fileHandleForReading.readabilityHandler = { handle in + let chunk = handle.availableData + if chunk.isEmpty { + handle.readabilityHandler = nil + group.leave() + return + } + queue.async { + collected.append(chunk) + } + } + + return { + group.wait() + return queue.sync { collected } + } +}