diff --git a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/MockDesugar.java b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/MockDesugar.java
index 5bd74cd596cd..d9ce5313f26e 100644
--- a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/MockDesugar.java
+++ b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/MockDesugar.java
@@ -115,35 +115,13 @@ public void generateMockFunctions(BLangPackage pkgNode) {
         // Get the Mock Function map from the pkgNode
         Map<String, String> mockFunctionMap = pkgNode.getTestablePkg().getMockFunctionNamesMap();
 
-        // Get all the imports symbols from the testable package
-        Set<String> mockFunctionSet = mockFunctionMap.keySet();
-        ArrayList<String> importsList = new ArrayList<>();
-        for (BLangImportPackage importPkg : pkgNode.getTestablePkg().getImports()) {
-            if (importPkg.symbol == null) {
-                continue;
-            }
-            if (!importPkg.symbol.toString().contains(testPackageSymbol)) {
-                importsList.add(importPkg.symbol.toString());
-            }
-        }
-
-        // Get all the imports from the current package
-        for (BLangImportPackage importPkg : pkgNode.getImports()) {
-            if (importPkg.symbol == null) {
-                continue;
-            }
-            if (importsList.contains(importPkg.symbol.toString())) {
-                continue;
-            }
-            if (!importPkg.symbol.toString().contains(testPackageSymbol)) {
-                importsList.add(importPkg.symbol.toString());
-            }
-        }
+        // Get the mock function type map from the pkgNode
+        Map<String, Boolean> isLegacyMockingMap = pkgNode.getTestablePkg().getIsLegacyMockingMap();
 
+        // Get the set of functions to generate
+        Set<String> mockFunctionSet = mockFunctionMap.keySet();
         for (String function : mockFunctionSet) {
-            if (function.contains(pkgNode.packageID.toString()) ? !function.split(pkgNode.packageID.toString())[1].
-                    startsWith(MOCK_LEGACY_DELIMITER) :
-                    !startsWithMockLegacyDelimiterForImportedMockFunctions(function, importsList)) {
+            if (!isLegacyMockingMap.get(function)) {
                 pkgNode.getTestablePkg().functions.add(generateMockFunction(function));
             }
         }
diff --git a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/tree/BLangTestablePackage.java b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/tree/BLangTestablePackage.java
index 412743a7b36c..87c9d4a9c3fa 100644
--- a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/tree/BLangTestablePackage.java
+++ b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/tree/BLangTestablePackage.java
@@ -29,6 +29,8 @@ public class BLangTestablePackage extends BLangPackage {
     // Semantic Data
     //Map to maintain all the mock functions
     private Map<String, String> mockFunctionNamesMap = new HashMap<>();
+
+    private final Map<String, Boolean> isLegacyMockingMap = new HashMap<>();
     public Map<String, String> getMockFunctionNamesMap() {
         return mockFunctionNamesMap;
     }
@@ -46,4 +48,12 @@ public <T> void accept(BLangNodeAnalyzer<T> analyzer, T props) {
     public <T, R> R apply(BLangNodeTransformer<T, R> modifier, T props) {
         return modifier.transform(this, props);
     }
+
+    public Map<String, Boolean> getIsLegacyMockingMap() {
+        return isLegacyMockingMap;
+    }
+
+    public void addIsLegacyMockingMap(String id, Boolean isLegacy) {
+        this.isLegacyMockingMap.put(id, isLegacy);
+    }
 }
diff --git a/misc/testerina/modules/testerina-core/src/main/java/org/ballerinalang/testerina/core/MockAnnotationProcessor.java b/misc/testerina/modules/testerina-core/src/main/java/org/ballerinalang/testerina/core/MockAnnotationProcessor.java
index 4535ce9d2173..1a60b64447c9 100644
--- a/misc/testerina/modules/testerina-core/src/main/java/org/ballerinalang/testerina/core/MockAnnotationProcessor.java
+++ b/misc/testerina/modules/testerina-core/src/main/java/org/ballerinalang/testerina/core/MockAnnotationProcessor.java
@@ -149,6 +149,8 @@ public void process(SimpleVariableNode simpleVariableNode, List<AnnotationAttach
                     bLangTestablePackage.addMockFunction(
                             functionToMockID + MOCK_FN_DELIMITER + annotationValues[1],
                             mockFnObjectName);
+                    bLangTestablePackage.addIsLegacyMockingMap(
+                            functionToMockID + MOCK_FN_DELIMITER + annotationValues[1], false);
 
                     if (functionToMockID != null) {
                         // Adding `<className> # <functionToMock> --> <MockFnObjectName>` to registry
@@ -257,6 +259,8 @@ public void process(FunctionNode functionNode, List<AnnotationAttachmentNode> an
                             (BLangTestablePackage) ((BLangFunction) functionNode).parent; // parent -> BLangPackage
                     bLangTestablePackage.addMockFunction(functionToMockID + MOCK_LEGACY_DELIMITER + vals[1],
                             functionName);
+                    bLangTestablePackage.addIsLegacyMockingMap(functionToMockID + MOCK_LEGACY_DELIMITER + vals[1],
+                            true);
 
                     // Adding `<className> # <functionToMock> --> <MockFnObjectName>` to registry
                     String className = getQualifiedClassName(bLangTestablePackage,