diff --git a/src/BlueWest.MapTo/EfAddGeneratorContext.cs b/src/BlueWest.MapTo/EfAddGeneratorContext.cs
index 195099d..c95357f 100644
--- a/src/BlueWest.MapTo/EfAddGeneratorContext.cs
+++ b/src/BlueWest.MapTo/EfAddGeneratorContext.cs
@@ -87,9 +87,9 @@ namespace MapTo
             }
         }
 
-        protected ImmutableArray<INamedTypeSymbol> GetEntityTypeSymbol(MemberDeclarationSyntax memberDeclarationSyntax, SemanticModel? semanticModel = null)
+        protected ImmutableArray<INamedTypeSymbol> GetEntityTypeSymbol(MemberDeclarationSyntax memberDeclarationSyntax, string attributeName, SemanticModel? semanticModel = null)
         {
-            var attributeData = memberDeclarationSyntax.GetAttribute(EfAddMethodsAttributeSource.AttributeName);
+            var attributeData = memberDeclarationSyntax.GetAttribute(attributeName);
             var sourceSymbol = GetEntityTypeSymbols(attributeData, semanticModel);
             return sourceSymbol;
         }
@@ -156,7 +156,7 @@ namespace MapTo
                 return null;
             }*/
             
-            var attributeTypeSymbols = GetEntityTypeSymbol(MemberSyntax, semanticModel);
+            var efAddAttributeTypeSymbols = GetEntityTypeSymbol(MemberSyntax, EfAddMethodsAttributeSource.AttributeName, semanticModel);
             
             // Main Entity
             var entityData = GetEntityTypeData(MemberSyntax, semanticModel);
@@ -166,17 +166,17 @@ namespace MapTo
             string entityTypeFullName = entityData.ToDisplayString(), createTypeFullName = entityData.ToDisplayString(), readTypeFullName = entityData.ToDisplayString();
 
 
-            if (attributeTypeSymbols.Length > 0)
+            if (efAddAttributeTypeSymbols.Length > 0)
             {
                 // Create DTO
-                createTypeIdentifierName = attributeTypeSymbols[0].Name;
-                createTypeFullName = attributeTypeSymbols[0].ToDisplayString();
+                createTypeIdentifierName = efAddAttributeTypeSymbols[0].Name;
+                createTypeFullName = efAddAttributeTypeSymbols[0].ToDisplayString();
             }
-            if (attributeTypeSymbols.Length > 1)
+            if (efAddAttributeTypeSymbols.Length > 1)
             {
                 // Read DTO
-                readTypeIdentifierName = attributeTypeSymbols[1].Name;
-                readTypeFullName = attributeTypeSymbols[1].ToDisplayString();
+                readTypeIdentifierName = efAddAttributeTypeSymbols[1].Name;
+                readTypeFullName = efAddAttributeTypeSymbols[1].ToDisplayString();
             }
             
             
@@ -204,7 +204,7 @@ namespace MapTo
             }
             
             
-            if (attributeTypeSymbols.IsDefaultOrEmpty)
+            if (efAddAttributeTypeSymbols.IsDefaultOrEmpty)
             {
                 AddDiagnostic(DiagnosticsFactory.MapFromAttributeNotFoundError(MemberSyntax.GetLocation()));
                 return null;
diff --git a/src/BlueWest.MapTo/Sources/EfMethodsSource.cs b/src/BlueWest.MapTo/Sources/EfMethodsSource.cs
index 11a25d4..cca6b49 100644
--- a/src/BlueWest.MapTo/Sources/EfMethodsSource.cs
+++ b/src/BlueWest.MapTo/Sources/EfMethodsSource.cs
@@ -26,17 +26,25 @@ namespace MapTo.Sources
             var propertyName = model.PropertyName;
             var newEntityVarName = $"new{model.EntityTypeIdentifierName}";
             var toCreateVarName = $"{model.EntityTypeIdentifierName.ToCamelCase()}ToCreate";
-
+            
+            builder
+                .WriteUsings(model.Usings)
+                .WriteLine("using Microsoft.EntityFrameworkCore;")
+                .WriteLine()
+                // Namespace declaration
+                .WriteLine($"namespace {model.Namespace}")
+                .WriteOpeningBracket()
+                // Class declaration
+                .WriteLine($"public static partial class {model.ContextTypeName}Extensions")
+                .WriteOpeningBracket()
+                .WriteLine();
+            
             if (!addSourceTemplate.IsEmpty())
             {
-                builder
-                    .WriteLine(GeneratedFilesHeader)
-                    .WriteUsings(model.Usings);
-                
+
                 var templateToSourceBuilder = new StringBuilder(addSourceTemplate);
                 templateToSourceBuilder
                         
-                    .Replace("{Namespace} ", model.Namespace)
                     .Replace("{entityTypeName}", entityTypeName)
                     .Replace("{readTypeFullName}", readTypeFullName)
                     .Replace("{createTypeFullName}", createTypeFullName)
@@ -55,16 +63,7 @@ namespace MapTo.Sources
             {
                 builder
                 .WriteLine(GeneratedFilesHeader)
-                    .WriteUsings(model.Usings)
-                    .WriteLine("using Microsoft.EntityFrameworkCore;")
-                    .WriteLine()
-                    // Namespace declaration
-                    .WriteLine($"namespace {model.Namespace}")
-                    .WriteOpeningBracket()
-                    // Class declaration
-                    .WriteLine($"public static partial class {model.EntityTypeIdentifierName}Extensions")
-                    .WriteOpeningBracket()
-                    .WriteLine()
+
                     .WriteLine($"public static (bool, {model.ReadTypeFullName}) Add{entityTypeName}(")
                     .WriteLine($"this {contextFullName} dbContext,")
                     .WriteLine($"{createTypeFullName} {toCreateVarName})")