diff --git a/src/main/java/Main.java b/src/main/java/Main.java
index 3fbd121..1f42d64 100644
--- a/src/main/java/Main.java
+++ b/src/main/java/Main.java
@@ -1,6 +1,7 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import designPatternExtensions.Adapter;
 import designPatternExtensions.DependencyInversion;
 import designPatternExtensions.InterfaceNode;
 import designPatternExtensions.MediatorInsertion;
@@ -31,14 +32,16 @@
 		try {
 			design = new DataTransferDesign(context);
 
-			ObjectNode cm = new ObjectNode("CompaniesMediator");
-			MediatorInsertion mi = new MediatorInsertion(customer, companies, cm);
+			ObjectNode cm = new ObjectNode("CompanyMediator");
+			MediatorInsertion mi = new MediatorInsertion(customer, company, cm);
 			design.addMediatorInsertion(mi);
 
-			DependencyInversion di = new DependencyInversion(company, new InterfaceNode(company, "ICompany"));
+			DependencyInversion di = new DependencyInversion(company, new InterfaceNode(company, "ICompany"), customer);
 			design.addDependencyInversion(di);
 
-			System.out.println(ASTGenerator.generate(design));
+
+			System.out.println(ASTGenerator.generate(Adapter.create()));
+//			System.out.println(ASTGenerator.generate(design));
 		} catch (IllegalRelationException e) {
 			e.printStackTrace();
 		}
diff --git a/src/main/java/designPatternExtensions/Adapter.java b/src/main/java/designPatternExtensions/Adapter.java
new file mode 100644
index 0000000..fe83a8c
--- /dev/null
+++ b/src/main/java/designPatternExtensions/Adapter.java
@@ -0,0 +1,38 @@
+package designPatternExtensions;
+
+import models.objectOrientedTransfer.*;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class Adapter {
+
+    public static DataTransferDesign create() throws IllegalRelationException {
+        ObjectNode main = new ObjectNode("Main");
+        ObjectNode fishingBoat = new ObjectNode("FishingBoat");
+        ObjectNode captain     = new ObjectNode("Captain");
+        ObjectNode data        = new ObjectNode("String");
+
+        List<ReferenceEdge> srcEdges = new ArrayList<>();
+        List<ReferenceEdge> dstEdges = new ArrayList<>();
+        ReferenceEdge srcEdge = new ReferenceEdge(main, fishingBoat, "fishingBoat", "String");
+        ReferenceEdge dstEdge = new ReferenceEdge(main, captain, "captain", "String");
+        srcEdges.add(srcEdge);
+        dstEdges.add(dstEdge);
+
+        List<Relation> relations = new ArrayList<>();
+        Relation transferRelation = new Relation(srcEdges, dstEdges, MultiplicityValue.OneToOne);
+        DataTransferContext context = new DataTransferContext(relations, transferRelation, data, "data", PushPullValue.PULL);
+
+        DataTransferDesign design = new DataTransferDesign(context);
+
+        ObjectNode adapter = new ObjectNode("FishingBoatAdapter");
+        MediatorInsertion mi = new MediatorInsertion(captain, fishingBoat, adapter);
+        design.addMediatorInsertion(mi);
+
+        DependencyInversion di = new DependencyInversion(adapter, new InterfaceNode(adapter, "RowingBoat"), captain);
+        design.addDependencyInversion(di);
+
+        return design;
+    }
+}
\ No newline at end of file
diff --git a/src/main/java/generators/ASTGenerator.java b/src/main/java/generators/ASTGenerator.java
index 7ae3eae..2eabbb0 100644
--- a/src/main/java/generators/ASTGenerator.java
+++ b/src/main/java/generators/ASTGenerator.java
@@ -310,6 +310,26 @@
 				replaceFieldNameInStatement(statement, dstFieldName, mediatorFieldName);
 			}
 		}
+
+		// Replace return type and MapType value type matching dst type with Mediator type in all classes
+		for (CompilationUnit cu: codebase.getCompilationUnits()) {
+			TypeDeclaration cls = cu.types().getFirst();
+			for (MethodDeclaration method: cls.getMethods()) {
+				if (!method.isConstructor() && method.getReturnType() != null) {
+					if (method.getReturnType().getTypeName().equals(dstType.getTypeName())) {
+						method.setReturnType(mediatorType);
+					}
+				}
+			}
+			for (FieldDeclaration field: cls.getFields()) {
+				if (field.getType() instanceof MapType) {
+					MapType mapType = (MapType) field.getType();
+					if (mapType.getValueType() != null && mapType.getValueType().getTypeName().equals(dstType.getTypeName())) {
+						field.setType(codebase.getMapType(mapType.getKeyType().getTypeName(), mediatorName));
+					}
+				}
+			}
+		}
 		return codebase;
 	}