From ff7d8c1930216d0049b69a3590220bb08b98c23b Mon Sep 17 00:00:00 2001
From: Ben Langmuir <blangmuir@apple.com>
Date: Thu, 22 Aug 2019 14:53:50 -0700
Subject: [PATCH] [test] Add listenToUnitEvents to initialization options

This should be used along with polling to explicitly control when index
data is read in tests that are outside the sourcekit-lsp process.
---
 Sources/SourceKit/SourceKitServer.swift | 14 ++++++++++++--
 Sources/SourceKit/Workspace.swift       | 17 +++++++++++++++--
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/Sources/SourceKit/SourceKitServer.swift b/Sources/SourceKit/SourceKitServer.swift
index d5359fc3..c8cdd807 100644
--- a/Sources/SourceKit/SourceKitServer.swift
+++ b/Sources/SourceKit/SourceKitServer.swift
@@ -222,18 +222,28 @@ extension SourceKitServer {
   // MARK: - General
 
   func initialize(_ req: Request<InitializeRequest>) {
+
+    var indexOptions = IndexOptions()
+    if case .dictionary(let options) = req.params.initializationOptions {
+      if case .bool(let listenToUnitEvents) = options["listenToUnitEvents"] {
+        indexOptions.listenToUnitEvents = listenToUnitEvents
+      }
+    }
+
     if let url = req.params.rootURL {
       self.workspace = try? Workspace(
         url: url,
         clientCapabilities: req.params.capabilities,
         toolchainRegistry: self.toolchainRegistry,
-        buildSetup: self.buildSetup)
+        buildSetup: self.buildSetup,
+        indexOptions: indexOptions)
     } else if let path = req.params.rootPath {
       self.workspace = try? Workspace(
         url: URL(fileURLWithPath: path),
         clientCapabilities: req.params.capabilities,
         toolchainRegistry: self.toolchainRegistry,
-        buildSetup: self.buildSetup)
+        buildSetup: self.buildSetup,
+        indexOptions: indexOptions)
     }
 
     if self.workspace == nil {
diff --git a/Sources/SourceKit/Workspace.swift b/Sources/SourceKit/Workspace.swift
index 11cbf129..fe176ed0 100644
--- a/Sources/SourceKit/Workspace.swift
+++ b/Sources/SourceKit/Workspace.swift
@@ -71,7 +71,8 @@ public final class Workspace {
     url: URL,
     clientCapabilities: ClientCapabilities,
     toolchainRegistry: ToolchainRegistry,
-    buildSetup: BuildSetup
+    buildSetup: BuildSetup,
+    indexOptions: IndexOptions = IndexOptions()
   ) throws {
 
     self.buildSetup = buildSetup
@@ -99,7 +100,8 @@ public final class Workspace {
         self.index = try IndexStoreDB(
           storePath: storePath.pathString,
           databasePath: dbPath.pathString,
-          library: lib)
+          library: lib,
+          listenToUnitEvents: indexOptions.listenToUnitEvents)
         log("opened IndexStoreDB at \(dbPath) with store path \(storePath)")
       } catch {
         log("failed to open IndexStoreDB: \(error.localizedDescription)", level: .error)
@@ -107,3 +109,14 @@ public final class Workspace {
     }
   }
 }
+
+public struct IndexOptions {
+
+  /// *For Testing* Whether the index should listen to unit events, or wait for
+  /// explicit calls to pollForUnitChangesAndWait().
+  public var listenToUnitEvents: Bool
+
+  public init(listenToUnitEvents: Bool = false) {
+    self.listenToUnitEvents = listenToUnitEvents
+  }
+}