Adds ability to add @testable import statements

Adds the ability to add @testable import statements via the -i flag. This is useful for users who want to only add their mocks to a project's Test target and use the @testable import functionality to import the main projects target

Author: maxwellE

README.md

@@ -95,6 +95,7 @@ OPTIONS:
   --annotation, -a        A custom annotation string used to indicate if a type should be mocked (default = @mockable).
   --header, -h            A custom header documentation to be added to the beginning of a generated mock file.
   --macro, -m             If set, #if [macro] / #endif will be added to the generated mock file content to guard compilation.
+  --testable-imports, -i  If set, @testable import statments will be added for each module name in this list.
   --concurrency-limit, -j Maximum number of threads to execute concurrently (default = number of cores on the running machine).
   --logging-level, -v     The logging level to use. Default is set to 0 (info only). Set 1 for verbose, 2 for warning, and 3 for error.
   --use-sourcekit         If this argument is added, it will use SourceKit for parsing. By default it uses SwiftSyntax.

Sources/Mockolo/Executor.swift

@@ -32,6 +32,7 @@ class Executor {
     private var exclusionSuffixes: OptionArgument<[String]>!
     private var header: OptionArgument<String>!
     private var macro: OptionArgument<String>!
+    private var testableImports: OptionArgument<[String]>!
     private var annotation: OptionArgument<String>!
     private var concurrencyLimit: OptionArgument<Int>!
     private var useSourceKit: OptionArgument<Bool>!
@@ -97,6 +98,10 @@ class Executor {
                                 shortName: "-m",
                                 kind: String.self,
                                 usage: "If set, #if [macro] / #endif will be added to the generated mock file content to guard compilation.")
+        testableImports = parser.add(option: "--testable-imports",
+                                        shortName: "-i",
+                                        kind: [String].self,
+                                        usage: "If set, @testable import statments will be added for each module name in this list.")
         header = parser.add(option: "--header",
                                 shortName: "-h",
                                 kind: String.self,
@@ -158,6 +163,7 @@ class Executor {
         let header = arguments.get(self.header)
         let loggingLevel = arguments.get(self.loggingLevel) ?? 0
         let macro = arguments.get(self.macro)
+        let testableImports = arguments.get(self.testableImports)
         let shouldUseSourceKit = arguments.get(useSourceKit) ?? false
 
         do {
@@ -169,6 +175,7 @@ class Executor {
                          annotation: annotation,
                          header: header,
                          macro: macro,
+                         testableImports: testableImports,
                          to: outputFilePath,
                          loggingLevel: loggingLevel,
                          concurrencyLimit: concurrencyLimit,

Sources/MockoloFramework/Operations/Generator.swift

@@ -36,6 +36,7 @@ public func generate(sourceDirs: [String]?,
                      annotation: String,
                      header: String?,
                      macro: String?,
+                     testableImports: [String]?,
                      to outputFilePath: String,
                      loggingLevel: Int,
                      concurrencyLimit: Int?,
@@ -160,6 +161,7 @@ public func generate(sourceDirs: [String]?,
                        pathToContentMap: pathToContentMap,
                        header: header,
                        macro: macro,
+                       testableImports: testableImports,
                        to: outputFilePath)
     signpost_end(name: "Write results")
     let t5 = CFAbsoluteTimeGetCurrent()

Sources/MockoloFramework/Operations/OutputWriter.swift

@@ -23,6 +23,7 @@ func write(candidates: [(String, Int64)],
            pathToContentMap: [(String, Data, Int64)],
            header: String?,
            macro: String?,
+           testableImports: [String]?,
            to outputFilePath: String) -> String {
 
     var importLines = [String]()
@@ -36,9 +37,27 @@ func write(candidates: [(String, Int64)],
         importLines.append(contentsOf: v)
         break
     }
-
-    let importsSet = Set(importLines.map{$0.trimmingCharacters(in: .whitespaces)})
-    let importLineStr = importsSet.sorted().joined(separator: "\n")
+    
+    var importLineStr = ""
+    
+    if let testableImports = testableImports {
+        var imports = importLines.compactMap { (importLine) -> String? in
+            return importLine.moduleName
+        }
+        imports.append(contentsOf: testableImports)
+        importLineStr = Set(imports)
+            .sorted()
+            .map { testableModuleName -> String in
+            guard testableImports.contains(testableModuleName) else {
+                return testableModuleName.asImport
+            }
+            return testableModuleName.asTestableImport
+        }
+        .joined(separator: "\n")
+    } else {
+        let importsSet = Set(importLines.map{$0.trimmingCharacters(in: .whitespaces)})
+        importLineStr = importsSet.sorted().joined(separator: "\n")
+    }
 
     let entities = candidates
         .sorted { (left: (String, Int64), right: (String, Int64)) -> Bool in

Sources/MockoloFramework/Utils/InheritanceResolver.swift

@@ -188,4 +188,3 @@ func findImportLines(data: Data, offset: Int64?) -> [String] {
 
     return []
 }
-

Sources/MockoloFramework/Utils/StringExtensions.swift

@@ -129,4 +129,16 @@ extension StringProtocol {
         return ret.components(separatedBy: separatorsForDisplay)
     }
 
+    var asTestableImport: String {
+        return "@testable \(self.asImport)"
+    }
+
+    var asImport: String {
+        return "import \(self)"
+    }
+    
+    var moduleName: String? {
+        guard self.hasPrefix(String.import) else { return nil }
+        return self.dropFirst(String.import.count).trimmingCharacters(in: CharacterSet.whitespaces)
+    }
 }

Tests/MockoloTestCase.swift

@@ -54,18 +54,18 @@ class MockoloTestCase: XCTestCase {
         }
     }
 
-    func verify(srcContent: String, mockContent: String? = nil, dstContent: String, header: String = "", concurrencyLimit: Int? = 1, parser: ParserType = .random) {
+    func verify(srcContent: String, mockContent: String? = nil, dstContent: String, header: String = "", testableImports: [String]? = [], concurrencyLimit: Int? = 1, parser: ParserType = .random) {
         var mockList: [String]?
         if let mock = mockContent {
             if mockList == nil {
                 mockList = [String]()
             }
             mockList?.append(mock)
         }
-        verify(srcContents: [srcContent], mockContents: mockList, dstContent: dstContent, header: header, concurrencyLimit: concurrencyLimit, parser: parser)
+        verify(srcContents: [srcContent], mockContents: mockList, dstContent: dstContent, header: header, testableImports: testableImports, concurrencyLimit: concurrencyLimit, parser: parser)
     }
 
-    func verify(srcContents: [String], mockContents: [String]?, dstContent: String, header: String, concurrencyLimit: Int?, parser: ParserType) {
+    func verify(srcContents: [String], mockContents: [String]?, dstContent: String, header: String, testableImports: [String]?, concurrencyLimit: Int?, parser: ParserType) {
         var index = 0
         srcFilePathsCount = srcContents.count
         mockFilePathsCount = mockContents?.count ?? 0
@@ -124,6 +124,7 @@ class MockoloTestCase: XCTestCase {
                       annotation: String.mockAnnotation,
                       header: header,
                       macro: "MOCK",
+                      testableImports: testableImports,
                       to: dstFilePath,
                       loggingLevel: 3,
                       concurrencyLimit: concurrencyLimit,

Tests/TestTestableImportStatements/FixtureTestableImportStatements.swift

@@ -0,0 +1,79 @@
+import MockoloFramework
+
+let testableImports = """
+\(String.headerDoc)
+import Foundation
+
+/// \(String.mockAnnotation)
+protocol SimpleVar {
+    var name: Int { get set }
+}
+"""
+
+let testableImportsMock =
+"""
+
+import Foundation
+@testable import SomeImport1
+@testable import SomeImport2
+
+
+class SimpleVarMock: SimpleVar {
+    
+    private var _doneInit = false
+    init() { _doneInit = true }
+    init(name: Int = 0) {
+        self.name = name
+        _doneInit = true
+    }
+    
+    var nameSetCallCount = 0
+    var underlyingName: Int = 0
+    var name: Int {
+        get { return underlyingName }
+        set {
+            underlyingName = newValue
+            if _doneInit { nameSetCallCount += 1 }
+        }
+    }
+}
+"""
+
+let testableImportsWithOverlap = """
+\(String.headerDoc)
+import Foundation
+import SomeImport1
+
+/// \(String.mockAnnotation)
+protocol SimpleVar {
+    var name: Int { get set }
+}
+"""
+
+let testableImportsWithOverlapMock =
+"""
+
+import Foundation
+@testable import SomeImport1
+
+
+class SimpleVarMock: SimpleVar {
+    
+    private var _doneInit = false
+    init() { _doneInit = true }
+    init(name: Int = 0) {
+        self.name = name
+        _doneInit = true
+    }
+    
+    var nameSetCallCount = 0
+    var underlyingName: Int = 0
+    var name: Int {
+        get { return underlyingName }
+        set {
+            underlyingName = newValue
+            if _doneInit { nameSetCallCount += 1 }
+        }
+    }
+}
+"""

Tests/TestTestableImportStatements/TestableImportStatementsTests.swift

@@ -0,0 +1,16 @@
+import Foundation
+
+class TestableImportStatementsTests: MockoloTestCase {
+    
+    func testTesableImportStatements() {
+        verify(srcContent: testableImports,
+               dstContent: testableImportsMock,
+               testableImports: ["SomeImport1", "SomeImport2"])
+    }
+
+    func testTesableImportStatementsWithOverlap() {
+        verify(srcContent: testableImportsWithOverlap,
+               dstContent: testableImportsWithOverlapMock,
+               testableImports: ["SomeImport1"])
+    }
+}