Fix Name Override With Inheritance

Sources/MockoloFramework/Operations/Generator.swift

@@ -53,6 +53,7 @@ public func generate(sourceDirs: [String],
     var candidates = [(String, Int64)]()
     var resolvedEntities = [ResolvedEntity]()
     var parentMocks = [String: Entity]()
+    var parentMocksByInheritedType = [String: Entity]()
     var protocolMap = [String: Entity]()
     var annotatedProtocolMap = [String: Entity]()
     var pathToImportsMap = ImportMap()
@@ -74,6 +75,12 @@ public func generate(sourceDirs: [String],
                                     }
         }
     }
+    // Build table mapping protocol name (inheritedType) => mock entity (entity)
+    for entity in parentMocks.values {
+        for inheritedType in entity.entityNode.inheritedTypes {
+            parentMocksByInheritedType[inheritedType] = entity
+        }
+    }
     signpost_end(name: "Process input")
     let t1 = CFAbsoluteTimeGetCurrent()
     log("Took", t1-t0, level: .verbose)
@@ -104,27 +111,33 @@ public func generate(sourceDirs: [String],
     let t2 = CFAbsoluteTimeGetCurrent()
     log("Took", t2-t1, level: .verbose)
 
-    let typeKeyList = [
-        parentMocks.compactMap { (key, value) -> String? in
-            if value.entityNode.mayHaveGlobalActor {
-                return nil
+    let typeKeyList = parentMocks.compactMap { (className, value) -> (String, String)? in
+        if value.entityNode.mayHaveGlobalActor {
+            return nil
+        }
+        let protocolName = className.components(separatedBy: "Mock").first ?? className
+        return (protocolName, "\(className)()")
+    } + annotatedProtocolMap
+        .lazy
+        .filter { _, entity in !entity.entityNode.mayHaveGlobalActor }
+        .map { typeName, entity in
+            switch entity.metadata?.nameOverride {
+            case .none:
+                (typeName, "\(typeName)Mock()")
+            case .some(let nameOverride):
+                (typeName, "\(nameOverride)()")
             }
-            return key.components(separatedBy: "Mock").first
-        },
-        annotatedProtocolMap.filter { !$0.value.entityNode.mayHaveGlobalActor }.map(\.key)
-    ]
-        .flatMap { $0 }
-        .map { typeName in
-            // nameOverride does not work correctly but it giving up.
-            return (typeName, "\(typeName)Mock()")
+
         }
+
     SwiftType.customDefaultValueMap = [String: String](typeKeyList, uniquingKeysWith: { $1 })
 
     signpost_begin(name: "Generate models")
     log("Resolve inheritance and generate unique entity models...", level: .info)
     generateUniqueModels(protocolMap: protocolMap,
                          annotatedProtocolMap: annotatedProtocolMap,
                          inheritanceMap: parentMocks,
+                         inheritanceByProtocolMap: parentMocksByInheritedType,
                          completion: { container in
                             resolvedEntities.append(container.entity)
                             relevantPaths.append(contentsOf: container.paths)

Sources/MockoloFramework/Operations/UniqueModelGenerator.swift

@@ -21,9 +21,10 @@ import Algorithms
 func generateUniqueModels(protocolMap: [String: Entity],
                           annotatedProtocolMap: [String: Entity],
                           inheritanceMap: [String: Entity],
+                          inheritanceByProtocolMap: [String: Entity],
                           completion: @escaping (ResolvedEntityContainer) -> ()) {
     scan(annotatedProtocolMap) { (key, val, lock) in
-        let ret = generateUniqueModels(key: key, entity: val, protocolMap: protocolMap, inheritanceMap: inheritanceMap)
+        let ret = generateUniqueModels(key: key, entity: val, protocolMap: protocolMap, inheritanceMap: inheritanceMap, inheritanceByProtocolMap: inheritanceByProtocolMap)
         lock?.lock()
         completion(ret)
         lock?.unlock()
@@ -33,8 +34,9 @@ func generateUniqueModels(protocolMap: [String: Entity],
 private func generateUniqueModels(key: String,
                                   entity: Entity,
                                   protocolMap: [String: Entity],
-                                  inheritanceMap: [String: Entity]) -> ResolvedEntityContainer {
-    let (models, processedModels, attributes, inheritedTypes, paths) = lookupEntities(key: key, declKind: entity.entityNode.declKind, protocolMap: protocolMap, inheritanceMap: inheritanceMap)
+                                  inheritanceMap: [String: Entity],
+                                  inheritanceByProtocolMap: [String: Entity]) -> ResolvedEntityContainer {
+    let (models, processedModels, attributes, inheritedTypes, paths) = lookupEntities(key: key, declKind: entity.entityNode.declKind, protocolMap: protocolMap, inheritanceMap: inheritanceMap, inheritanceByProtocolMap: inheritanceByProtocolMap)
 
     let processedFullNames = processedModels.compactMap {$0.fullName}
 

Sources/MockoloFramework/Parsers/SourceParser.swift

@@ -32,7 +32,7 @@ class SourceParser {
                              fileMacro: String?,
                              completion: @escaping ([Entity], ImportMap?) -> ()) {
         scan(paths) { (path, lock) in
-            self.generateASTs(path, annotation: "", fileMacro: fileMacro, declType: .classType, lock: lock, completion: completion)
+            self.generateASTs(path, annotation: "", fileMacro: fileMacro, declType: .classType, classesHaveBeenProcessed: true, lock: lock, completion: completion)
         }
     }
     /// Parses decls (protocol, class) with annotations (/// @mockable) and calls a completion block
@@ -68,6 +68,7 @@ class SourceParser {
                               annotation: String,
                               fileMacro: String?,
                               declType: FindTargetDeclType,
+                              classesHaveBeenProcessed: Bool = false,
                               lock: NSLock?,
                               completion: @escaping ([Entity], ImportMap?) -> ()) {
 
@@ -87,7 +88,7 @@ class SourceParser {
 
         var results = [Entity]()
         let node = Parser.parse(path)
-        let treeVisitor = EntityVisitor(path, annotation: annotation, fileMacro: fileMacro, declType: declType)
+        let treeVisitor = EntityVisitor(path, annotation: annotation, fileMacro: fileMacro, declType: declType, classesHaveBeenProcessed: classesHaveBeenProcessed)
         treeVisitor.walk(node)
         let ret = treeVisitor.entities
         results.append(contentsOf: ret)

Sources/MockoloFramework/Parsers/SwiftSyntaxExtensions.swift

@@ -724,11 +724,13 @@ final class EntityVisitor: SyntaxVisitor {
     let fileMacro: String
     let path: String
     let declType: FindTargetDeclType
-    init(_ path: String, annotation: String = "", fileMacro: String?, declType: FindTargetDeclType) {
+    let classesHaveBeenProcessed: Bool
+    init(_ path: String, annotation: String = "", fileMacro: String?, declType: FindTargetDeclType, classesHaveBeenProcessed: Bool = false) {
         self.annotation = annotation
         self.fileMacro = fileMacro ?? ""
         self.path = path
         self.declType = declType
+        self.classesHaveBeenProcessed = classesHaveBeenProcessed
         super.init(viewMode: .sourceAccurate)
     }
 
@@ -749,7 +751,7 @@ final class EntityVisitor: SyntaxVisitor {
     }
 
     override func visit(_ node: ClassDeclSyntax) -> SyntaxVisitorContinueKind {
-        if node.nameText.hasSuffix("Mock") {
+        if classesHaveBeenProcessed || node.nameText.hasSuffix("Mock") {
             // this mock class node must be public else wouldn't have compiled before
             if let ent = Entity.node(with: node, filepath: path, isPrivate: node.isPrivate, isFinal: false, metadata: nil, processed: true) {
                 entities.append(ent)

Sources/MockoloFramework/Utils/InheritanceResolver.swift

@@ -22,12 +22,14 @@ import Algorithms
 /// @param key The entity name to look up
 /// @param protocolMap Used to look up the current entity and its inheritance types
 /// @param inheritanceMap Used to look up inherited types if not contained in protocolMap
+/// @param inheritanceByProtocolMap Used to look up names of mock implementation for protocols
 /// @returns a list of models representing sub-entities of the current entity, a list of models processed in dependent mock files if exists,
 ///          cumulated attributes, cumulated inherited types, and a map of filepaths and file contents (used for import lines lookup later).
 func lookupEntities(key: String,
                     declKind: NominalTypeDeclKind,
                     protocolMap: [String: Entity],
-                    inheritanceMap: [String: Entity]) -> ([Model], [Model], [String], Set<String>, [String]) {
+                    inheritanceMap: [String: Entity],
+                    inheritanceByProtocolMap: [String: Entity]) -> ([Model], [Model], [String], Set<String>, [String]) {
 
     // Used to keep track of types to be mocked
     var models = [Model]()
@@ -55,7 +57,7 @@ func lookupEntities(key: String,
             // If the protocol inherits other protocols, look up their entities as well.
             for parent in current.entityNode.inheritedTypes {
                 if parent != .class, parent != .anyType, parent != .anyObject {
-                    let (parentModels, parentProcessedModels, parentAttributes, parentInheritedTypes, parentPaths) = lookupEntities(key: parent, declKind: declKind, protocolMap: protocolMap, inheritanceMap: inheritanceMap)
+                    let (parentModels, parentProcessedModels, parentAttributes, parentInheritedTypes, parentPaths) = lookupEntities(key: parent, declKind: declKind, protocolMap: protocolMap, inheritanceMap: inheritanceMap, inheritanceByProtocolMap: inheritanceByProtocolMap)
                     models.append(contentsOf: parentModels)
                     processedModels.append(contentsOf: parentProcessedModels)
                     attributes.append(contentsOf: parentAttributes)
@@ -64,7 +66,7 @@ func lookupEntities(key: String,
                 }
             }
         }
-    } else if let parentMock = inheritanceMap["\(key)Mock"], declKind == .protocol {
+    } else if let parentMock = inheritanceMap["\(key)Mock"] ?? inheritanceByProtocolMap[key], declKind == .protocol {
         // If the parent protocol is not in the protocol map, look it up in the input parent mocks map.
         let sub = parentMock.entityNode.subContainer(metadata: parentMock.metadata, declKind: declKind, path: parentMock.filepath, isProcessed: parentMock.isProcessed)
         processedModels.append(contentsOf: sub.members)

Tests/TestOverrides/FixtureMockNames.swift

@@ -23,3 +23,92 @@ class FooMock: FooProtocol {
     }
 }
 """
+
+func mockableDeclaration(customName: String?) -> String {
+    switch customName {
+    case .some(let customName):
+        "/// @mockable(override: name = \(customName))"
+    case .none:
+        "/// @mockable()"
+    }
+}
+
+func baseProtocol(customName: String?) -> String {
+    """
+    \(mockableDeclaration(customName: customName))
+    protocol BaseProtocol {
+        func register()
+        var counter: Int { get } 
+    }
+    """
+}
+
+func baseMock(customName: String?) -> String {
+    """
+    class \(customName ?? "BaseProtocolMock"): BaseProtocol {
+        init() { }
+        init(counter: Int = 0) {
+            self.counter = counter
+        }
+        private(set) var registerCallCount = 0
+        var registerHandler: (() -> ())?
+        func register() {
+            registerCallCount += 1
+            if let registerHandler = registerHandler {
+                registerHandler()
+            }
+        }
+        var counter: Int = 0
+    }
+    """
+}
+
+func derivedProtocol(customName: String?) -> String {
+    """
+    \(mockableDeclaration(customName: customName))
+    protocol DerivedProtocol : BaseProtocol {
+        func like()
+        func subscribe() 
+    }
+    """
+}
+
+func derivedMock(customName: String?) -> String {
+    """
+    class \(customName ?? "DerivedProtocolMock"): DerivedProtocol {
+        init() { }
+        init(counter: Int = 0) {
+            self.counter = counter
+        }
+
+        private(set) var likeCallCount = 0
+        var likeHandler: (() -> ())?
+        func like() {
+            likeCallCount += 1
+            if let likeHandler = likeHandler {
+                likeHandler()
+            }
+
+        }
+
+        private(set) var subscribeCallCount = 0
+        var subscribeHandler: (() -> ())?
+        func subscribe() {
+            subscribeCallCount += 1
+            if let subscribeHandler = subscribeHandler {
+                subscribeHandler()
+            }
+
+        }
+        private(set) var registerCallCount = 0
+        var registerHandler: (() -> ())?
+        func register() {
+            registerCallCount += 1
+            if let registerHandler = registerHandler {
+                registerHandler()
+            }
+        }
+        var counter: Int = 0
+    }
+    """
+}

Tests/TestOverrides/NameOverrideTests.swift

@@ -4,4 +4,27 @@ class NameOverrideTests: MockoloTestCase {
     func testNameOverride() {
         verify(srcContent: nameOverride, dstContent: nameOverrideMock)
     }
+
+    let baseCustomizations: [String?] = [nil, "BaseProtocolMock", "BaseMock", "FakeBase"]
+    let derivedCustomizations: [String?] = [nil, "DerivedProtocolMock", "DerivedMock", "FakeDerived"]
+
+    func testBaseFixtures() {
+        for baseCustomization in baseCustomizations {
+            verify(srcContent: baseProtocol(customName: baseCustomization), dstContent: baseMock(customName: baseCustomization))
+        }
+    }
+
+    func testDerivedFixtures() {
+        for baseCustomization in baseCustomizations {
+            let baseImplementation = baseMock(customName: baseCustomization)
+            for derivedCustomization in derivedCustomizations {
+                verify(
+                    srcContent: derivedProtocol(customName: derivedCustomization),
+                    mockContent: baseImplementation,
+                    dstContent: derivedMock(customName: derivedCustomization)
+                )
+            }
+        }
+    }
+
 }