diff --git a/resources/fileTemplates/internal/Magento Delete Entity By Id Command.php.ft b/resources/fileTemplates/internal/Magento Delete Entity By Id Command.php.ft
index ac18b5b21..625d76ff0 100644
--- a/resources/fileTemplates/internal/Magento Delete Entity By Id Command.php.ft	
+++ b/resources/fileTemplates/internal/Magento Delete Entity By Id Command.php.ft	
@@ -56,9 +56,9 @@ class ${CLASS_NAME}
         try {
             /** @var ${MODEL} $model */
             $model = $this->modelFactory->create();
-            $this->resource->load($model, $entityId, '${ENTITY_ID}');
+            $this->resource->load($model, $entityId, ${ENTITY_ID_REFERENCE});
 
-            if (!$model->getData('${ENTITY_ID}')) {
+            if (!$model->getData(${ENTITY_ID_REFERENCE})) {
                 throw new ${NO_SUCH_ENTITY_EXCEPTION}(
                     __('Could not find ${ENTITY_NAME} with id: `%id`',
                         [
diff --git a/resources/fileTemplates/internal/Magento Entity Delete Controller Class.php.ft b/resources/fileTemplates/internal/Magento Entity Delete Controller Class.php.ft
index 01097476c..8c83d6c23 100644
--- a/resources/fileTemplates/internal/Magento Entity Delete Controller Class.php.ft	
+++ b/resources/fileTemplates/internal/Magento Entity Delete Controller Class.php.ft	
@@ -47,7 +47,7 @@ class ${CLASS_NAME} extends ${EXTENDS} implements ${IMPLEMENTS_POST}, ${IMPLEMEN
         /** @var ${RESULT_INTERFACE} $resultRedirect */
         $resultRedirect = $this->resultFactory->create(${RESULT_FACTORY}::TYPE_REDIRECT);
         $resultRedirect->setPath('*/*/');
-        $entityId = (int) $this->getRequest()->getParam('${ENTITY_ID}');
+        $entityId = (int) $this->getRequest()->getParam(${ENTITY_ID_REFERENCE});
 
         try {
             $this->deleteByIdCommand->execute($entityId);
diff --git a/resources/fileTemplates/internal/Magento Entity Save Controller Class.php.ft b/resources/fileTemplates/internal/Magento Entity Save Controller Class.php.ft
index 552eac78c..d15b064b4 100644
--- a/resources/fileTemplates/internal/Magento Entity Save Controller Class.php.ft	
+++ b/resources/fileTemplates/internal/Magento Entity Save Controller Class.php.ft	
@@ -77,7 +77,7 @@ class ${CLASS_NAME} extends ${EXTENDS} implements ${IMPLEMENTS}
                 $this->dataPersistor->set('entity', $params);
 
                 return $resultRedirect->setPath('*/*/edit', [
-                '${ENTITY_ID}'=> $this->getRequest()->getParam('${ENTITY_ID}')
+                ${ENTITY_ID_REFERENCE} => $this->getRequest()->getParam(${ENTITY_ID_REFERENCE})
                 ]);
         }
 
diff --git a/resources/fileTemplates/internal/Magento Grid Ui Component Action Column Class.php.ft b/resources/fileTemplates/internal/Magento Grid Ui Component Action Column Class.php.ft
index 06a29a555..5f2757a9c 100644
--- a/resources/fileTemplates/internal/Magento Grid Ui Component Action Column Class.php.ft	
+++ b/resources/fileTemplates/internal/Magento Grid Ui Component Action Column Class.php.ft	
@@ -59,12 +59,12 @@ class ${CLASS_NAME} extends ${PARENT_CLASS}
      */
     public function prepareDataSource(array $dataSource): array
     {
-    #set($entityIdAccessor = "['" + ${ENTITY_ID} + "']")
+    #set($entityIdAccessor = "[" + ${ENTITY_ID_REFERENCE} + "]")
         if (isset($dataSource['data']['items'])) {
             foreach ($dataSource['data']['items'] as &$item) {
                 if (isset($item$entityIdAccessor)) {
                     $entityName = static::ENTITY_NAME;
-                    $urlData = ['${ENTITY_ID}' => $item$entityIdAccessor];
+                    $urlData = [${ENTITY_ID_REFERENCE} => $item$entityIdAccessor];
 
                     $editUrl = $this->urlBuilder->getUrl(static::EDIT_URL_PATH, $urlData);
                     $deleteUrl = $this->urlBuilder->getUrl(static::DELETE_URL_PATH, $urlData);
diff --git a/resources/fileTemplates/internal/Magento PHP Form Generic Button Block Class.php.ft b/resources/fileTemplates/internal/Magento PHP Form Generic Button Block Class.php.ft
index 998ebb28a..aafe31dc4 100644
--- a/resources/fileTemplates/internal/Magento PHP Form Generic Button Block Class.php.ft	
+++ b/resources/fileTemplates/internal/Magento PHP Form Generic Button Block Class.php.ft	
@@ -40,7 +40,7 @@ class ${CLASS_NAME}
      */
     public function ${ENTITY_ID_GETTER}(): int
     {
-        return (int) $this->context->getRequest()->getParam('${ENTITY_ID}');
+        return (int) $this->context->getRequest()->getParam(${ENTITY_ID_REFERENCE});
     }
 
     /**
diff --git a/resources/fileTemplates/internal/Magento Resource Model Class.php.ft b/resources/fileTemplates/internal/Magento Resource Model Class.php.ft
index bb33568bb..f2a3d580d 100644
--- a/resources/fileTemplates/internal/Magento Resource Model Class.php.ft	
+++ b/resources/fileTemplates/internal/Magento Resource Model Class.php.ft	
@@ -21,7 +21,11 @@ class ${NAME}#if (${EXTENDS}) extends ${EXTENDS}#end#if (${IMPLEMENTS}) implemen
      */
     protected function _construct()
     {
+#if (${ENTITY_ID_REFERENCE})
+        $this->_init('${DB_NAME}', ${ENTITY_ID_REFERENCE});
+#else
         $this->_init('${DB_NAME}', '${ENTITY_ID_COLUMN}');
+#end
         $this->_useIsObjectNew = true;
     }
 }
diff --git a/resources/fileTemplates/internal/Magento Save Entity Command Model.php.ft b/resources/fileTemplates/internal/Magento Save Entity Command Model.php.ft
index 39bf47820..cd2e50e3a 100644
--- a/resources/fileTemplates/internal/Magento Save Entity Command Model.php.ft	
+++ b/resources/fileTemplates/internal/Magento Save Entity Command Model.php.ft	
@@ -59,7 +59,7 @@ class ${CLASS_NAME}
             $model->addData($${DTO_PROPERTY}->getData());
             $model->setHasDataChanges(true);
 
-            if (!$model->getId()) {
+            if (!$model->getData(${ENTITY_ID_CONST})) {
                 $model->isObjectNew(true);
             }
             $this->resource->save($model);
diff --git a/resources/fileTemplates/internal/Magento UI Component Custom Data Provider Class.php.ft b/resources/fileTemplates/internal/Magento UI Component Custom Data Provider Class.php.ft
index 610889926..183e69f4f 100644
--- a/resources/fileTemplates/internal/Magento UI Component Custom Data Provider Class.php.ft	
+++ b/resources/fileTemplates/internal/Magento UI Component Custom Data Provider Class.php.ft	
@@ -87,7 +87,7 @@ class ${CLASS_NAME} extends ${EXTENDS}
             $result->getItems(),
             $result->getTotalCount(),
             $searchCriteria,
-            '#if(${ENTITY_ID})${ENTITY_ID}#{else}entity_id#end'
+            ${ENTITY_ID_REFERENCE}
         );
     }
 
@@ -104,14 +104,14 @@ class ${CLASS_NAME} extends ${EXTENDS}
         $this->loadedData = parent::getData();
         $itemsById = [];
 
-#set($entityId = "#if(${ENTITY_ID})${ENTITY_ID}#{else}entity_id#end")
-#set($entityIdAccessor = '[(int) $item[' + "'$entityId'" + ']]')
+#set($entityId = "${ENTITY_ID_REFERENCE}")
+#set($entityIdAccessor = '[(int) $item[' + "${ENTITY_ID_REFERENCE}" + ']]')
 #set($getIdAccessor = "[(int) $id]")
         foreach ($this->loadedData['items'] as $item) {
             $itemsById$entityIdAccessor = $item;
         }
 
-        if ($id = $this->request->getParam('$entityId', null)) {
+        if ($id = $this->request->getParam($entityId)) {
             $this->loadedData['entity'] = $itemsById$getIdAccessor;
         }
 
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/context/EntityCreatorContext.java b/src/com/magento/idea/magento2plugin/actions/generation/context/EntityCreatorContext.java
new file mode 100644
index 000000000..8383a1649
--- /dev/null
+++ b/src/com/magento/idea/magento2plugin/actions/generation/context/EntityCreatorContext.java
@@ -0,0 +1,15 @@
+/*
+ * Copyright © Magento, Inc. All rights reserved.
+ * See COPYING.txt for license details.
+ */
+
+package com.magento.idea.magento2plugin.actions.generation.context;
+
+import com.intellij.openapi.util.Key;
+import com.intellij.openapi.util.UserDataHolderBase;
+
+public class EntityCreatorContext extends UserDataHolderBase implements GenerationContext {
+
+    public static final Key<String> DTO_TYPE = Key.create("DTO_TYPE");
+    public static final Key<String> ENTITY_ID = Key.create("ENTITY_ID");
+}
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/context/GenerationContext.java b/src/com/magento/idea/magento2plugin/actions/generation/context/GenerationContext.java
new file mode 100644
index 000000000..58f5f9888
--- /dev/null
+++ b/src/com/magento/idea/magento2plugin/actions/generation/context/GenerationContext.java
@@ -0,0 +1,9 @@
+/*
+ * Copyright © Magento, Inc. All rights reserved.
+ * See COPYING.txt for license details.
+ */
+
+package com.magento.idea.magento2plugin.actions.generation.context;
+
+public interface GenerationContext {
+}
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/data/SaveEntityCommandData.java b/src/com/magento/idea/magento2plugin/actions/generation/data/SaveEntityCommandData.java
index ecea167c0..e84fe479b 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/data/SaveEntityCommandData.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/data/SaveEntityCommandData.java
@@ -11,6 +11,7 @@ public class SaveEntityCommandData {
 
     private final String moduleName;
     private final String entityName;
+    private final String entityId;
     private final String modelName;
     private final String resourceModelName;
     private final String dtoName;
@@ -23,6 +24,7 @@ public class SaveEntityCommandData {
      *
      * @param moduleName String
      * @param entityName String
+     * @param entityId String
      * @param modelName String
      * @param resourceModelName String
      * @param dtoName String
@@ -33,6 +35,7 @@ public class SaveEntityCommandData {
     public SaveEntityCommandData(
             final @NotNull String moduleName,
             final @NotNull String entityName,
+            final @NotNull String entityId,
             final @NotNull String modelName,
             final @NotNull String resourceModelName,
             final @NotNull String dtoName,
@@ -42,6 +45,7 @@ public SaveEntityCommandData(
     ) {
         this.moduleName = moduleName;
         this.entityName = entityName;
+        this.entityId = entityId;
         this.modelName = modelName;
         this.resourceModelName = resourceModelName;
         this.dtoName = dtoName;
@@ -68,6 +72,15 @@ public String getEntityName() {
         return entityName;
     }
 
+    /**
+     * Get entity id.
+     *
+     * @return String
+     */
+    public String getEntityId() {
+        return entityId;
+    }
+
     /**
      * Get model name.
      *
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/data/converter/newentitydialog/SaveEntityCommandDtoConverter.java b/src/com/magento/idea/magento2plugin/actions/generation/data/converter/newentitydialog/SaveEntityCommandDtoConverter.java
index bc2b0086c..8704e8e16 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/data/converter/newentitydialog/SaveEntityCommandDtoConverter.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/data/converter/newentitydialog/SaveEntityCommandDtoConverter.java
@@ -27,6 +27,7 @@ public SaveEntityCommandDtoConverter(
         super(
                 generationContextData.getModuleName(),
                 newEntityDialogData.getEntityName(),
+                newEntityDialogData.getIdFieldName(),
                 newEntityDialogData.getEntityName().concat("Model"),
                 newEntityDialogData.getEntityName().concat("Resource"),
                 newEntityDialogData.getEntityName().concat("Data"),
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/dialog/NewEntityDialog.java b/src/com/magento/idea/magento2plugin/actions/generation/dialog/NewEntityDialog.java
index 38db2f6a0..b1d6d724b 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/dialog/NewEntityDialog.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/dialog/NewEntityDialog.java
@@ -10,6 +10,7 @@
 import com.intellij.ui.DocumentAdapter;
 import com.intellij.util.indexing.FileBasedIndex;
 import com.magento.idea.magento2plugin.actions.generation.NewEntityAction;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.UiComponentFormButtonData;
 import com.magento.idea.magento2plugin.actions.generation.data.UiComponentFormFieldData;
 import com.magento.idea.magento2plugin.actions.generation.data.UiComponentFormFieldsetData;
@@ -36,6 +37,7 @@
 import com.magento.idea.magento2plugin.actions.generation.generator.pool.provider.NewEntityGeneratorsProviderUtil;
 import com.magento.idea.magento2plugin.actions.generation.generator.util.DbSchemaGeneratorUtil;
 import com.magento.idea.magento2plugin.actions.generation.generator.util.NamespaceBuilder;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 import com.magento.idea.magento2plugin.magento.files.ControllerBackendPhp;
 import com.magento.idea.magento2plugin.magento.files.DataModelFile;
 import com.magento.idea.magento2plugin.magento.files.DataModelInterfaceFile;
@@ -404,6 +406,15 @@ private void onOK() {
 
         final NewEntityDialogData dialogData = getNewEntityDialogData();
         final EntityCreatorContextData context = getEntityCreatorContextData(dialogData);
+        final EntityCreatorContext generationContext = new EntityCreatorContext();
+        generationContext.putUserData(
+                EntityCreatorContext.DTO_TYPE,
+                dialogData.hasDtoInterface()
+                        ? context.getDtoInterfaceNamespaceBuilder().getClassFqn()
+                        : context.getDtoModelNamespaceBuilder().getClassFqn()
+        );
+        generationContext.putUserData(EntityCreatorContext.ENTITY_ID, dialogData.getIdFieldName());
+        GenerationContextRegistry.getInstance().setContext(generationContext);
 
         final GeneratorPoolHandler generatorPoolHandler = new GeneratorPoolHandler(context);
 
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/dialog/util/ClassPropertyFormatterUtil.java b/src/com/magento/idea/magento2plugin/actions/generation/dialog/util/ClassPropertyFormatterUtil.java
index 3c1a4363e..08ab105c0 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/dialog/util/ClassPropertyFormatterUtil.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/dialog/util/ClassPropertyFormatterUtil.java
@@ -7,6 +7,7 @@
 
 import com.google.common.base.CaseFormat;
 import com.magento.idea.magento2plugin.actions.generation.data.code.ClassPropertyData;
+import com.magento.idea.magento2plugin.actions.generation.generator.util.PhpClassGeneratorUtil;
 import java.util.ArrayList;
 import java.util.List;
 import javax.swing.table.DefaultTableModel;
@@ -50,7 +51,7 @@ public static String formatSingleProperty(final String name, final String type)
                 CaseFormat.LOWER_UNDERSCORE.to(CaseFormat.LOWER_CAMEL, name),
                 CaseFormat.LOWER_UNDERSCORE.to(CaseFormat.UPPER_CAMEL, name),
                 name,
-                CaseFormat.LOWER_UNDERSCORE.to(CaseFormat.UPPER_UNDERSCORE, name)
+                formatNameToConstant(name)
         ).string();
     }
 
@@ -61,4 +62,30 @@ public static String formatSingleProperty(final String name, final String type)
     public static String joinProperties(final List<String> properties) {
         return StringUtils.join(properties, ",");
     }
+
+    /**
+     * Format property field to constant format with specified type.
+     *
+     * @param name String
+     *
+     * @return String
+     */
+    public static String formatNameToConstant(final String name, final String typeFqn) {
+        return String.format(
+                "%s::%s",
+                PhpClassGeneratorUtil.getNameFromFqn(typeFqn),
+                formatNameToConstant(name)
+        );
+    }
+
+    /**
+     * Format property field to constant format.
+     *
+     * @param name String
+     *
+     * @return String
+     */
+    public static String formatNameToConstant(final String name) {
+        return CaseFormat.LOWER_UNDERSCORE.to(CaseFormat.UPPER_UNDERSCORE, name);
+    }
 }
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityByIdCommandGenerator.java b/src/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityByIdCommandGenerator.java
index 24ada3467..23b2aaafd 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityByIdCommandGenerator.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityByIdCommandGenerator.java
@@ -6,13 +6,17 @@
 package com.magento.idea.magento2plugin.actions.generation.generator;
 
 import com.intellij.openapi.project.Project;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.DeleteEntityByIdCommandData;
+import com.magento.idea.magento2plugin.actions.generation.dialog.util.ClassPropertyFormatterUtil;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 import com.magento.idea.magento2plugin.magento.files.AbstractPhpFile;
 import com.magento.idea.magento2plugin.magento.files.ModelFile;
 import com.magento.idea.magento2plugin.magento.files.ResourceModelFile;
 import com.magento.idea.magento2plugin.magento.files.commands.DeleteEntityByIdCommandFile;
 import com.magento.idea.magento2plugin.magento.packages.code.ExceptionType;
 import com.magento.idea.magento2plugin.magento.packages.code.FrameworkLibraryType;
+import java.util.Objects;
 import java.util.Properties;
 import org.jetbrains.annotations.NotNull;
 
@@ -67,11 +71,22 @@ protected void fillAttributes(final @NotNull Properties attributes) {
         final ResourceModelFile resourceFile =
                 new ResourceModelFile(data.getModuleName(), data.getResourceModelName());
 
+        final EntityCreatorContext context =
+                (EntityCreatorContext) GenerationContextRegistry.getInstance().getContext();
+        Objects.requireNonNull(context);
+        final String dtoTypeFqn = context.getUserData(EntityCreatorContext.DTO_TYPE);
+        Objects.requireNonNull(dtoTypeFqn);
+        typesBuilder.append(
+                "ENTITY_ID_REFERENCE",
+                ClassPropertyFormatterUtil.formatNameToConstant(data.getEntityId(), dtoTypeFqn),
+                false
+        );
+        typesBuilder.append("DTO_TYPE", dtoTypeFqn);
+
         typesBuilder
                 .append("ENTITY_NAME", data.getEntityName(), false)
                 .append("NAMESPACE", file.getNamespace(), false)
                 .append("CLASS_NAME", DeleteEntityByIdCommandFile.CLASS_NAME, false)
-                .append("ENTITY_ID", data.getEntityId(), false)
                 .append("Exception", "Exception")
                 .append("COULD_NOT_DELETE", ExceptionType.COULD_NOT_DELETE.getType())
                 .append("NO_SUCH_ENTITY_EXCEPTION",
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityControllerFileGenerator.java b/src/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityControllerFileGenerator.java
index d8e7bff6d..6e5d62da4 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityControllerFileGenerator.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityControllerFileGenerator.java
@@ -6,7 +6,10 @@
 package com.magento.idea.magento2plugin.actions.generation.generator;
 
 import com.intellij.openapi.project.Project;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.DeleteEntityControllerFileData;
+import com.magento.idea.magento2plugin.actions.generation.dialog.util.ClassPropertyFormatterUtil;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 import com.magento.idea.magento2plugin.magento.files.AbstractPhpFile;
 import com.magento.idea.magento2plugin.magento.files.actions.DeleteActionFile;
 import com.magento.idea.magento2plugin.magento.files.commands.DeleteEntityByIdCommandFile;
@@ -14,6 +17,7 @@
 import com.magento.idea.magento2plugin.magento.packages.code.BackendModuleType;
 import com.magento.idea.magento2plugin.magento.packages.code.ExceptionType;
 import com.magento.idea.magento2plugin.magento.packages.code.FrameworkLibraryType;
+import java.util.Objects;
 import java.util.Properties;
 import org.jetbrains.annotations.NotNull;
 
@@ -67,7 +71,6 @@ protected void fillAttributes(final @NotNull Properties attributes) {
                 .append("ENTITY_NAME", data.getEntityName(), false)
                 .append("CLASS_NAME", DeleteActionFile.CLASS_NAME, false)
                 .append("ADMIN_RESOURCE", data.getAcl(), false)
-                .append("ENTITY_ID", data.getEntityId(), false)
                 .append("DELETE_COMMAND",
                         new DeleteEntityByIdCommandFile(
                                 data.getModuleName(),
@@ -84,5 +87,17 @@ protected void fillAttributes(final @NotNull Properties attributes) {
                 .append("NO_SUCH_ENTITY_EXCEPTION",
                         ExceptionType.NO_SUCH_ENTITY_EXCEPTION.getType())
                 .append("COULD_NOT_DELETE", ExceptionType.COULD_NOT_DELETE.getType());
+
+        final EntityCreatorContext context =
+                (EntityCreatorContext) GenerationContextRegistry.getInstance().getContext();
+        Objects.requireNonNull(context);
+        final String dtoTypeFqn = context.getUserData(EntityCreatorContext.DTO_TYPE);
+        Objects.requireNonNull(dtoTypeFqn);
+        typesBuilder.append(
+                "ENTITY_ID_REFERENCE",
+                ClassPropertyFormatterUtil.formatNameToConstant(data.getEntityId(), dtoTypeFqn),
+                false
+        );
+        typesBuilder.append("DTO_TYPE", dtoTypeFqn);
     }
 }
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/generator/FormGenericButtonBlockGenerator.java b/src/com/magento/idea/magento2plugin/actions/generation/generator/FormGenericButtonBlockGenerator.java
index 64c875a13..e464cfb4c 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/generator/FormGenericButtonBlockGenerator.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/generator/FormGenericButtonBlockGenerator.java
@@ -6,12 +6,16 @@
 package com.magento.idea.magento2plugin.actions.generation.generator;
 
 import com.intellij.openapi.project.Project;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.FormGenericButtonBlockData;
+import com.magento.idea.magento2plugin.actions.generation.dialog.util.ClassPropertyFormatterUtil;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 import com.magento.idea.magento2plugin.magento.files.AbstractPhpFile;
 import com.magento.idea.magento2plugin.magento.files.FormGenericButtonBlockFile;
 import com.magento.idea.magento2plugin.magento.packages.code.FrameworkLibraryType;
 import java.util.Arrays;
 import java.util.Locale;
+import java.util.Objects;
 import java.util.Properties;
 import java.util.stream.Collectors;
 import org.jetbrains.annotations.NotNull;
@@ -65,11 +69,24 @@ protected void fillAttributes(final @NotNull Properties attributes) {
                 .map(s -> s.substring(0, 1).toUpperCase(Locale.getDefault()) + s.substring(1))
                 .collect(Collectors.joining());
 
+        final EntityCreatorContext context =
+                (EntityCreatorContext) GenerationContextRegistry.getInstance().getContext();
+
+        if (context != null) {
+            final String dtoTypeFqn = context.getUserData(EntityCreatorContext.DTO_TYPE);
+            Objects.requireNonNull(dtoTypeFqn);
+            typesBuilder.append(
+                    "ENTITY_ID_REFERENCE",
+                    ClassPropertyFormatterUtil.formatNameToConstant(data.getEntityId(), dtoTypeFqn),
+                    false
+            );
+            typesBuilder.append("DTO_TYPE", dtoTypeFqn);
+        }
+
         typesBuilder
                 .append("NAMESPACE", file.getNamespace(), false)
                 .append("ENTITY_NAME", data.getEntityName(), false)
                 .append("CLASS_NAME", FormGenericButtonBlockFile.CLASS_NAME, false)
-                .append("ENTITY_ID", data.getEntityId(), false)
                 .append("ENTITY_ID_GETTER", entityIdGetter, false)
                 .append("CONTEXT", FormGenericButtonBlockFile.CONTEXT)
                 .append("URL", FrameworkLibraryType.URL.getType());
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/generator/GridActionColumnFileGenerator.java b/src/com/magento/idea/magento2plugin/actions/generation/generator/GridActionColumnFileGenerator.java
index 795a7e7a0..cc2627652 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/generator/GridActionColumnFileGenerator.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/generator/GridActionColumnFileGenerator.java
@@ -6,10 +6,14 @@
 package com.magento.idea.magento2plugin.actions.generation.generator;
 
 import com.intellij.openapi.project.Project;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.GridActionColumnData;
+import com.magento.idea.magento2plugin.actions.generation.dialog.util.ClassPropertyFormatterUtil;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 import com.magento.idea.magento2plugin.magento.files.AbstractPhpFile;
 import com.magento.idea.magento2plugin.magento.files.GridActionColumnFile;
 import com.magento.idea.magento2plugin.magento.packages.code.FrameworkLibraryType;
+import java.util.Objects;
 import java.util.Properties;
 import org.jetbrains.annotations.NotNull;
 
@@ -60,7 +64,6 @@ protected AbstractPhpFile initFile() {
     protected void fillAttributes(final @NotNull Properties attributes) {
         typesBuilder
                 .append("ENTITY_NAME", data.getEntityName(), false)
-                .append("ENTITY_ID", data.getEntityIdColumn(), false)
                 .append("NAMESPACE", file.getNamespace(), false)
                 .append("CLASS_NAME", file.getClassName(), false)
                 .append("EDIT_URL_PATH", data.getEditUrlPath(), false)
@@ -69,5 +72,20 @@ protected void fillAttributes(final @NotNull Properties attributes) {
                 .append("URL", FrameworkLibraryType.URL.getType())
                 .append("CONTEXT", GridActionColumnFile.CONTEXT)
                 .append("UI_COMPONENT_FACTORY", GridActionColumnFile.UI_COMPONENT_FACTORY);
+
+        final EntityCreatorContext context =
+                (EntityCreatorContext) GenerationContextRegistry.getInstance().getContext();
+        Objects.requireNonNull(context);
+        final String dtoTypeFqn = context.getUserData(EntityCreatorContext.DTO_TYPE);
+        Objects.requireNonNull(dtoTypeFqn);
+        typesBuilder.append(
+                "ENTITY_ID_REFERENCE",
+                ClassPropertyFormatterUtil.formatNameToConstant(
+                        data.getEntityIdColumn(),
+                        dtoTypeFqn
+                ),
+                false
+        );
+        typesBuilder.append("DTO_TYPE", dtoTypeFqn);
     }
 }
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/generator/ModuleResourceModelGenerator.java b/src/com/magento/idea/magento2plugin/actions/generation/generator/ModuleResourceModelGenerator.java
index 877b2af40..24c525f2c 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/generator/ModuleResourceModelGenerator.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/generator/ModuleResourceModelGenerator.java
@@ -6,9 +6,13 @@
 package com.magento.idea.magento2plugin.actions.generation.generator;
 
 import com.intellij.openapi.project.Project;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.ResourceModelData;
+import com.magento.idea.magento2plugin.actions.generation.dialog.util.ClassPropertyFormatterUtil;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 import com.magento.idea.magento2plugin.magento.files.AbstractPhpFile;
 import com.magento.idea.magento2plugin.magento.files.ResourceModelFile;
+import java.util.Objects;
 import java.util.Properties;
 import org.jetbrains.annotations.NotNull;
 
@@ -63,5 +67,22 @@ protected void fillAttributes(final @NotNull Properties attributes) {
                 .append("DB_NAME", data.getDbTableName(), false)
                 .append("ENTITY_ID_COLUMN", data.getEntityIdColumn(), false)
                 .append("EXTENDS", ResourceModelFile.ABSTRACT_DB);
+
+        final EntityCreatorContext context =
+                (EntityCreatorContext) GenerationContextRegistry.getInstance().getContext();
+
+        if (context != null && context.getUserData(EntityCreatorContext.DTO_TYPE) != null) {
+            final String dtoTypeFqn = context.getUserData(EntityCreatorContext.DTO_TYPE);
+            Objects.requireNonNull(dtoTypeFqn);
+            typesBuilder.append(
+                    "ENTITY_ID_REFERENCE",
+                    ClassPropertyFormatterUtil.formatNameToConstant(
+                            data.getEntityIdColumn(),
+                            dtoTypeFqn
+                    ),
+                    false
+            );
+            typesBuilder.append("DTO_TYPE", dtoTypeFqn);
+        }
     }
 }
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityCommandGenerator.java b/src/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityCommandGenerator.java
index 301d46265..bb851d641 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityCommandGenerator.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityCommandGenerator.java
@@ -8,6 +8,7 @@
 import com.google.common.base.CaseFormat;
 import com.intellij.openapi.project.Project;
 import com.magento.idea.magento2plugin.actions.generation.data.SaveEntityCommandData;
+import com.magento.idea.magento2plugin.actions.generation.dialog.util.ClassPropertyFormatterUtil;
 import com.magento.idea.magento2plugin.magento.files.AbstractPhpFile;
 import com.magento.idea.magento2plugin.magento.files.DataModelFile;
 import com.magento.idea.magento2plugin.magento.files.DataModelInterfaceFile;
@@ -80,18 +81,23 @@ protected void fillAttributes(final @NotNull Properties attributes) {
         final String modelType = modelFile.getClassFqn();
         final String modelFactoryType = modelType.concat("Factory");
         final String resourceType = resourceModelFile.getClassFqn();
+        String dtoType;
 
         if (data.isDtoWithInterface()) {
             final DataModelInterfaceFile dataModelInterfaceFile =
                     new DataModelInterfaceFile(data.getModuleName(), data.getDtoInterfaceName());
-            final String dtoType = dataModelInterfaceFile.getClassFqn();
-            typesBuilder.append("DTO", dtoType);
+            dtoType = dataModelInterfaceFile.getClassFqn();
         } else {
             final DataModelFile dataModelFile =
                     new DataModelFile(data.getModuleName(), data.getDtoName());
-            final String dtoType = dataModelFile.getClassFqn();
-            typesBuilder.append("DTO", dtoType);
+            dtoType = dataModelFile.getClassFqn();
         }
+        typesBuilder.append("DTO", dtoType);
+        typesBuilder.append(
+                "ENTITY_ID_CONST",
+                ClassPropertyFormatterUtil.formatNameToConstant(data.getEntityId(), dtoType),
+                false
+        );
 
         final String dtoProperty = CaseFormat.UPPER_CAMEL.to(
                 CaseFormat.LOWER_CAMEL, data.getEntityName()
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityControllerFileGenerator.java b/src/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityControllerFileGenerator.java
index 5c853b6b9..cfa897d37 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityControllerFileGenerator.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityControllerFileGenerator.java
@@ -7,6 +7,7 @@
 
 import com.intellij.openapi.project.Project;
 import com.magento.idea.magento2plugin.actions.generation.data.SaveEntityControllerFileData;
+import com.magento.idea.magento2plugin.actions.generation.dialog.util.ClassPropertyFormatterUtil;
 import com.magento.idea.magento2plugin.magento.files.AbstractPhpFile;
 import com.magento.idea.magento2plugin.magento.files.DataModelFile;
 import com.magento.idea.magento2plugin.magento.files.DataModelInterfaceFile;
@@ -74,12 +75,16 @@ protected void fillAttributes(final @NotNull Properties attributes) {
                     new DataModelFile(data.getModuleName(), data.getDtoName());
             dtoType = dataModelFile.getClassFqn();
         }
+        typesBuilder.append(
+                "ENTITY_ID_REFERENCE",
+                ClassPropertyFormatterUtil.formatNameToConstant(data.getEntityId(), dtoType),
+                false
+        );
 
         typesBuilder
                 .append("NAMESPACE", file.getNamespace(), false)
                 .append("ENTITY_NAME", data.getEntityName(), false)
                 .append("CLASS_NAME", SaveActionFile.CLASS_NAME, false)
-                .append("ENTITY_ID", data.getEntityId(), false)
                 .append("ADMIN_RESOURCE", data.getAcl(), false)
                 .append("IMPLEMENTS", HttpMethod.POST.getInterfaceFqn())
                 .append("DATA_PERSISTOR", FrameworkLibraryType.DATA_PERSISTOR.getType())
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/generator/UiComponentDataProviderGenerator.java b/src/com/magento/idea/magento2plugin/actions/generation/generator/UiComponentDataProviderGenerator.java
index bea312eb6..d2f76344d 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/generator/UiComponentDataProviderGenerator.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/generator/UiComponentDataProviderGenerator.java
@@ -6,11 +6,15 @@
 package com.magento.idea.magento2plugin.actions.generation.generator;
 
 import com.intellij.openapi.project.Project;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.UiComponentDataProviderData;
+import com.magento.idea.magento2plugin.actions.generation.dialog.util.ClassPropertyFormatterUtil;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 import com.magento.idea.magento2plugin.magento.files.AbstractPhpFile;
 import com.magento.idea.magento2plugin.magento.files.UiComponentDataProviderFile;
 import com.magento.idea.magento2plugin.magento.files.queries.GetListQueryFile;
 import com.magento.idea.magento2plugin.magento.packages.code.FrameworkLibraryType;
+import java.util.Objects;
 import java.util.Properties;
 import org.jetbrains.annotations.NotNull;
 
@@ -72,7 +76,20 @@ protected void fillAttributes(final @NotNull Properties attributes) {
                 .append("EXTENDS", UiComponentDataProviderFile.DEFAULT_DATA_PROVIDER);
 
         if (data.getEntityIdFieldName() != null && data.getEntityName() != null) {
-            typesBuilder.append("ENTITY_ID", data.getEntityIdFieldName(), false);
+            final EntityCreatorContext context =
+                    (EntityCreatorContext) GenerationContextRegistry.getInstance().getContext();
+            Objects.requireNonNull(context);
+            final String dtoTypeFqn = context.getUserData(EntityCreatorContext.DTO_TYPE);
+            Objects.requireNonNull(dtoTypeFqn);
+            typesBuilder.append(
+                    "ENTITY_ID_REFERENCE",
+                    ClassPropertyFormatterUtil.formatNameToConstant(
+                            data.getEntityIdFieldName(),
+                            dtoTypeFqn
+                    ),
+                    false
+            );
+            typesBuilder.append("DTO_TYPE", dtoTypeFqn);
 
             typesBuilder
                     .append("HAS_GET_LIST_QUERY", "true", false)
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/generator/UiComponentFormButtonBlockGenerator.java b/src/com/magento/idea/magento2plugin/actions/generation/generator/UiComponentFormButtonBlockGenerator.java
index c525eef15..3015e34e8 100644
--- a/src/com/magento/idea/magento2plugin/actions/generation/generator/UiComponentFormButtonBlockGenerator.java
+++ b/src/com/magento/idea/magento2plugin/actions/generation/generator/UiComponentFormButtonBlockGenerator.java
@@ -9,11 +9,14 @@
 import com.intellij.psi.PsiDirectory;
 import com.intellij.psi.PsiFile;
 import com.jetbrains.php.lang.psi.elements.PhpClass;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.UiComponentFormButtonData;
+import com.magento.idea.magento2plugin.actions.generation.dialog.util.ClassPropertyFormatterUtil;
 import com.magento.idea.magento2plugin.actions.generation.generator.util.DirectoryGenerator;
 import com.magento.idea.magento2plugin.actions.generation.generator.util.FileFromTemplateGenerator;
 import com.magento.idea.magento2plugin.actions.generation.generator.util.NamespaceBuilder;
 import com.magento.idea.magento2plugin.actions.generation.generator.util.PhpClassGeneratorUtil;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 import com.magento.idea.magento2plugin.indexes.ModuleIndex;
 import com.magento.idea.magento2plugin.magento.files.FormButtonBlockFile;
 import com.magento.idea.magento2plugin.magento.files.FormGenericButtonBlockFile;
@@ -25,6 +28,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Properties;
 import java.util.stream.Collectors;
 import org.bouncycastle.util.Strings;
@@ -178,7 +182,21 @@ protected void fillAttributes(final @NotNull Properties attributes) {
 
         final Map<String, String> variables = new HashMap<>();
         variables.put("$varName", Strings.toLowerCase(entityName));
-        variables.put("$varId", entityIdField);
+
+        final EntityCreatorContext context =
+                (EntityCreatorContext) GenerationContextRegistry.getInstance().getContext();
+
+        if (buttonTypeSettings.isVariableExpected("$varIdConst") && context != null) {
+            final String dtoTypeFqn = context.getUserData(EntityCreatorContext.DTO_TYPE);
+            Objects.requireNonNull(dtoTypeFqn);
+            variables.put(
+                    "$varIdConst",
+                    ClassPropertyFormatterUtil.formatNameToConstant(entityIdField, dtoTypeFqn)
+            );
+            uses.add(dtoTypeFqn);
+        } else {
+            variables.put("$varIdConst", entityIdField);
+        }
         variables.put("$varEntityIdAccessor", entityIdAccessor);
 
         attributes.setProperty("ON_CLICK", buttonTypeSettings.getOnClick(variables));
diff --git a/src/com/magento/idea/magento2plugin/actions/generation/util/GenerationContextRegistry.java b/src/com/magento/idea/magento2plugin/actions/generation/util/GenerationContextRegistry.java
new file mode 100644
index 000000000..f8b29682e
--- /dev/null
+++ b/src/com/magento/idea/magento2plugin/actions/generation/util/GenerationContextRegistry.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright © Magento, Inc. All rights reserved.
+ * See COPYING.txt for license details.
+ */
+
+package com.magento.idea.magento2plugin.actions.generation.util;
+
+import com.magento.idea.magento2plugin.actions.generation.context.GenerationContext;
+
+/**
+ * This class used to share some generation context data through
+ * generation lifetime to not overweight data objects.
+ */
+public final class GenerationContextRegistry {
+
+    private static final GenerationContextRegistry INSTANCE = new GenerationContextRegistry();
+    private GenerationContext context;
+
+    private GenerationContextRegistry() {}
+
+    /**
+     * Get generation context registry.
+     *
+     * @return GenerationContextRegistry
+     */
+    public static GenerationContextRegistry getInstance() {
+        return INSTANCE;
+    }
+
+    /**
+     * Set context.
+     *
+     * @param context GenerationContext
+     */
+    public void setContext(final GenerationContext context) {
+        this.context = context;
+    }
+
+    /**
+     * Get context.
+     *
+     * @return GenerationContext
+     */
+    public GenerationContext getContext() {
+        return context;
+    }
+}
diff --git a/src/com/magento/idea/magento2plugin/magento/packages/UiFormButtonTypeSettings.java b/src/com/magento/idea/magento2plugin/magento/packages/UiFormButtonTypeSettings.java
index f0e23626e..3a7ab6cb0 100644
--- a/src/com/magento/idea/magento2plugin/magento/packages/UiFormButtonTypeSettings.java
+++ b/src/com/magento/idea/magento2plugin/magento/packages/UiFormButtonTypeSettings.java
@@ -17,8 +17,8 @@ public enum UiFormButtonTypeSettings {
             + "            ]", 10, "Save entity button."),
     DELETE("Delete", "delete", "'deleteConfirm(\\''\n"
             + "            . __('Are you sure you want to delete this $varName?')\n"
-            + "            . '\\', \\'' . $this->getUrl('*/*/delete', ['$varId' => "
-            + "$this->$varEntityIdAccessor]) . '\\')'", "[]", 20, "Delete entity button."),
+            + "            . '\\', \\'' . $this->getUrl(\n'*/*/delete',\n[$varIdConst => "
+            + "$this->$varEntityIdAccessor]\n) . '\\')'", "[]", 20, "Delete entity button."),
     BACK("Back To Grid", "back",
             "sprintf(\"location.href = '%s';\", $this->getUrl('*/*/'))",
             "[]", 30, "Back to list button."),
@@ -116,6 +116,17 @@ public String getDataAttrs(final Map<String, String> variables) {
         return parseVariables(dataAttrs, variables);
     }
 
+    /**
+     * Check if specified variable used for current button.
+     *
+     * @param variableName String
+     *
+     * @return boolean
+     */
+    public boolean isVariableExpected(final @NotNull String variableName) {
+        return onClick.contains(variableName) || dataAttrs.contains(variableName);
+    }
+
     /**
      * Get button sort order.
      *
diff --git a/testData/actions/generation/generator/DeleteEntityCommandGenerator/generateDeleteEntityByIdCommandFile/DeleteByIdCommand.php b/testData/actions/generation/generator/DeleteEntityCommandGenerator/generateDeleteEntityByIdCommandFile/DeleteByIdCommand.php
index 014032232..404852ee2 100644
--- a/testData/actions/generation/generator/DeleteEntityCommandGenerator/generateDeleteEntityByIdCommandFile/DeleteByIdCommand.php
+++ b/testData/actions/generation/generator/DeleteEntityCommandGenerator/generateDeleteEntityByIdCommandFile/DeleteByIdCommand.php
@@ -5,6 +5,7 @@
 use Exception;
 use Foo\Bar\Model\BookModel;
 use Foo\Bar\Model\BookModelFactory;
+use Foo\Bar\Model\Data\BookData;
 use Foo\Bar\Model\ResourceModel\BookResource;
 use Magento\Framework\Exception\CouldNotDeleteException;
 use Magento\Framework\Exception\NoSuchEntityException;
@@ -59,9 +60,9 @@ public function execute(int $entityId)
         try {
             /** @var BookModel $model */
             $model = $this->modelFactory->create();
-            $this->resource->load($model, $entityId, 'book_id');
+            $this->resource->load($model, $entityId, BookData::BOOK_ID);
 
-            if (!$model->getData('book_id')) {
+            if (!$model->getData(BookData::BOOK_ID)) {
                 throw new NoSuchEntityException(
                     __('Could not find Book with id: `%id`',
                         [
diff --git a/testData/actions/generation/generator/DeleteEntityGenerator/generateDeleteEntityFile/Delete.php b/testData/actions/generation/generator/DeleteEntityGenerator/generateDeleteEntityFile/Delete.php
index 279f1b597..92594d159 100644
--- a/testData/actions/generation/generator/DeleteEntityGenerator/generateDeleteEntityFile/Delete.php
+++ b/testData/actions/generation/generator/DeleteEntityGenerator/generateDeleteEntityFile/Delete.php
@@ -3,6 +3,7 @@
 namespace Foo\Bar\Controller\Adminhtml\Company;
 
 use Foo\Bar\Command\Company\DeleteByIdCommand;
+use Foo\Bar\Model\Data\CompanyData;
 use Magento\Backend\App\Action;
 use Magento\Backend\App\Action\Context;
 use Magento\Framework\App\Action\HttpGetActionInterface;
@@ -52,7 +53,7 @@ public function execute()
         /** @var ResultInterface $resultRedirect */
         $resultRedirect = $this->resultFactory->create(ResultFactory::TYPE_REDIRECT);
         $resultRedirect->setPath('*/*/');
-        $entityId = (int)$this->getRequest()->getParam('company_id');
+        $entityId = (int)$this->getRequest()->getParam(CompanyData::COMPANY_ID);
 
         try {
             $this->deleteByIdCommand->execute($entityId);
diff --git a/testData/actions/generation/generator/FormButtonBlockGenerator/generateDeleteButtonBlock/DeleteBlock.php b/testData/actions/generation/generator/FormButtonBlockGenerator/generateDeleteButtonBlock/DeleteBlock.php
index dc3814f81..b2d2ad729 100644
--- a/testData/actions/generation/generator/FormButtonBlockGenerator/generateDeleteButtonBlock/DeleteBlock.php
+++ b/testData/actions/generation/generator/FormButtonBlockGenerator/generateDeleteButtonBlock/DeleteBlock.php
@@ -2,6 +2,7 @@
 
 namespace Foo\Bar\Block\Form\Book;
 
+use Foo\Bar\Model\Data\BookData;
 use Magento\Framework\View\Element\UiComponent\Control\ButtonProviderInterface;
 
 /**
@@ -21,7 +22,10 @@ public function getButtonData(): array
             'delete',
             'deleteConfirm(\''
             . __('Are you sure you want to delete this book?')
-            . '\', \'' . $this->getUrl('*/*/delete', ['book_id' => $this->getBookId()]) . '\')',
+            . '\', \'' . $this->getUrl(
+                '*/*/delete',
+                [BookData::BOOK_ID => $this->getBookId()]
+            ) . '\')',
             [],
             20
         );
diff --git a/testData/actions/generation/generator/FormGenericButtonBlockGenerator/generateFormGenericButtonBlockFile/GenericButton.php b/testData/actions/generation/generator/FormGenericButtonBlockGenerator/generateFormGenericButtonBlockFile/GenericButton.php
index 00d30defb..18f8b2d7e 100644
--- a/testData/actions/generation/generator/FormGenericButtonBlockGenerator/generateFormGenericButtonBlockFile/GenericButton.php
+++ b/testData/actions/generation/generator/FormGenericButtonBlockGenerator/generateFormGenericButtonBlockFile/GenericButton.php
@@ -2,6 +2,7 @@
 
 namespace Foo\Bar\Block\Form\Book;
 
+use Foo\Bar\Model\Data\BookData;
 use Magento\Backend\Block\Widget\Context;
 use Magento\Framework\UrlInterface;
 
@@ -38,7 +39,7 @@ public function __construct(
      */
     public function getBookId(): int
     {
-        return (int)$this->context->getRequest()->getParam('book_id');
+        return (int)$this->context->getRequest()->getParam(BookData::BOOK_ID);
     }
 
     /**
diff --git a/testData/actions/generation/generator/GridActionColumnFileGenerator/generateGridActionColumnFile/BookBlockActions.php b/testData/actions/generation/generator/GridActionColumnFileGenerator/generateGridActionColumnFile/BookBlockActions.php
index 5b7bb0c6f..383555ecf 100644
--- a/testData/actions/generation/generator/GridActionColumnFileGenerator/generateGridActionColumnFile/BookBlockActions.php
+++ b/testData/actions/generation/generator/GridActionColumnFileGenerator/generateGridActionColumnFile/BookBlockActions.php
@@ -2,6 +2,7 @@
 
 namespace Foo\Bar\Ui\Component\Listing\Column;
 
+use Foo\Bar\Model\Data\BookData;
 use Magento\Framework\UrlInterface;
 use Magento\Framework\View\Element\UiComponentFactory;
 use Magento\Framework\View\Element\UiComponent\ContextInterface;
@@ -61,9 +62,9 @@ public function prepareDataSource(array $dataSource): array
     {
         if (isset($dataSource['data']['items'])) {
             foreach ($dataSource['data']['items'] as &$item) {
-                if (isset($item['book_id'])) {
+                if (isset($item[BookData::BOOK_ID])) {
                     $entityName = static::ENTITY_NAME;
-                    $urlData = ['book_id' => $item['book_id']];
+                    $urlData = [BookData::BOOK_ID => $item[BookData::BOOK_ID]];
 
                     $editUrl = $this->urlBuilder->getUrl(static::EDIT_URL_PATH, $urlData);
                     $deleteUrl = $this->urlBuilder->getUrl(static::DELETE_URL_PATH, $urlData);
diff --git a/testData/actions/generation/generator/ModuleResourceModelGenerator/generateFileWithDtoReference/TestResourceModel.php b/testData/actions/generation/generator/ModuleResourceModelGenerator/generateFileWithDtoReference/TestResourceModel.php
new file mode 100644
index 000000000..4eea307e0
--- /dev/null
+++ b/testData/actions/generation/generator/ModuleResourceModelGenerator/generateFileWithDtoReference/TestResourceModel.php
@@ -0,0 +1,23 @@
+<?php
+
+namespace Foo\Bar\Model\ResourceModel;
+
+use Foo\Bar\Model\Data\EntityData;
+use Magento\Framework\Model\ResourceModel\Db\AbstractDb;
+
+class TestResourceModel extends AbstractDb
+{
+    /**
+     * @var string
+     */
+    protected $_eventPrefix = 'my_table_resource_model';
+
+    /**
+     * @inheritdoc
+     */
+    protected function _construct()
+    {
+        $this->_init('my_table', EntityData::ENTITY_ID);
+        $this->_useIsObjectNew = true;
+    }
+}
diff --git a/testData/actions/generation/generator/SaveEntityActionGenerator/generateSaveEntityActionWithInterfaceFile/Save.php b/testData/actions/generation/generator/SaveEntityActionGenerator/generateSaveEntityActionWithInterfaceFile/Save.php
index dcff220fd..1780a3a7b 100644
--- a/testData/actions/generation/generator/SaveEntityActionGenerator/generateSaveEntityActionWithInterfaceFile/Save.php
+++ b/testData/actions/generation/generator/SaveEntityActionGenerator/generateSaveEntityActionWithInterfaceFile/Save.php
@@ -84,7 +84,7 @@ public function execute()
             $this->dataPersistor->set('entity', $params);
 
             return $resultRedirect->setPath('*/*/edit', [
-                'entity_id' => $this->getRequest()->getParam('entity_id')
+                CompanyInterface::COMPANY_ID => $this->getRequest()->getParam(CompanyInterface::COMPANY_ID)
             ]);
         }
 
diff --git a/testData/actions/generation/generator/SaveEntityActionGenerator/generateSaveEntityActionWithoutInterfaceFile/Save.php b/testData/actions/generation/generator/SaveEntityActionGenerator/generateSaveEntityActionWithoutInterfaceFile/Save.php
index 2bac73502..4875c60da 100644
--- a/testData/actions/generation/generator/SaveEntityActionGenerator/generateSaveEntityActionWithoutInterfaceFile/Save.php
+++ b/testData/actions/generation/generator/SaveEntityActionGenerator/generateSaveEntityActionWithoutInterfaceFile/Save.php
@@ -84,7 +84,7 @@ public function execute()
             $this->dataPersistor->set('entity', $params);
 
             return $resultRedirect->setPath('*/*/edit', [
-                'entity_id' => $this->getRequest()->getParam('entity_id')
+                CompanyData::COMPANY_ID => $this->getRequest()->getParam(CompanyData::COMPANY_ID)
             ]);
         }
 
diff --git a/testData/actions/generation/generator/SaveEntityCommandGenerator/generateSaveEntityCommandFile/SaveCommand.php b/testData/actions/generation/generator/SaveEntityCommandGenerator/generateSaveEntityCommandFile/SaveCommand.php
index 533eb48e5..8a28e9071 100644
--- a/testData/actions/generation/generator/SaveEntityCommandGenerator/generateSaveEntityCommandFile/SaveCommand.php
+++ b/testData/actions/generation/generator/SaveEntityCommandGenerator/generateSaveEntityCommandFile/SaveCommand.php
@@ -63,7 +63,7 @@ public function execute(BookData $book): int
             $model->addData($book->getData());
             $model->setHasDataChanges(true);
 
-            if (!$model->getId()) {
+            if (!$model->getData(BookData::BOOK_ID)) {
                 $model->isObjectNew(true);
             }
             $this->resource->save($model);
diff --git a/testData/actions/generation/generator/UiComponentGridDataProviderGenerator/generateDataProviderWithInjectedGetListQuery/GridDataProvider.php b/testData/actions/generation/generator/UiComponentGridDataProviderGenerator/generateDataProviderWithInjectedGetListQuery/GridDataProvider.php
index c9156caff..75837df65 100644
--- a/testData/actions/generation/generator/UiComponentGridDataProviderGenerator/generateDataProviderWithInjectedGetListQuery/GridDataProvider.php
+++ b/testData/actions/generation/generator/UiComponentGridDataProviderGenerator/generateDataProviderWithInjectedGetListQuery/GridDataProvider.php
@@ -2,6 +2,7 @@
 
 namespace Foo\Bar\Ui\Component\Listing;
 
+use Foo\Bar\Model\Data\BookData;
 use Foo\Bar\Query\Book\GetListQuery;
 use Magento\Framework\Api\FilterBuilder;
 use Magento\Framework\Api\Search\ReportingInterface;
@@ -84,7 +85,7 @@ public function getSearchResult()
             $result->getItems(),
             $result->getTotalCount(),
             $searchCriteria,
-            'entity_id'
+            BookData::BOOK_ID
         );
     }
 
@@ -102,10 +103,10 @@ public function getData(): array
         $itemsById = [];
 
         foreach ($this->loadedData['items'] as $item) {
-            $itemsById[(int)$item['entity_id']] = $item;
+            $itemsById[(int)$item[BookData::BOOK_ID]] = $item;
         }
 
-        if ($id = $this->request->getParam('entity_id', null)) {
+        if ($id = $this->request->getParam(BookData::BOOK_ID)) {
             $this->loadedData['entity'] = $itemsById[(int)$id];
         }
 
diff --git a/tests/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityCommandGeneratorTest.java b/tests/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityCommandGeneratorTest.java
index 8933d848d..c16fe7a78 100644
--- a/tests/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityCommandGeneratorTest.java
+++ b/tests/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityCommandGeneratorTest.java
@@ -6,7 +6,9 @@
 package com.magento.idea.magento2plugin.actions.generation.generator;
 
 import com.intellij.psi.PsiFile;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.DeleteEntityByIdCommandData;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 import com.magento.idea.magento2plugin.magento.files.commands.DeleteEntityByIdCommandFile;
 
 public class DeleteEntityCommandGeneratorTest extends BaseGeneratorTestCase {
@@ -15,10 +17,20 @@ public class DeleteEntityCommandGeneratorTest extends BaseGeneratorTestCase {
     private static final String ENTITY_NAME = "Book";
     private static final String MODEL_NAME = ENTITY_NAME + "Model";
     private static final String RESOURCE_MODEL_NAME = ENTITY_NAME + "Resource";
+    private static final String ENTITY_DTO_TYPE = "Foo\\Bar\\Model\\Data\\BookData";
     private static final String EXPECTED_DIRECTORY = "src/app/code/Foo/Bar/Command/" + ENTITY_NAME;
     private static final String ENTITY_ID = "book_id";
     private static final String ACL = "Foo_Bar::book_management";
 
+    @Override
+    protected void setUp() throws Exception {
+        super.setUp();
+        final EntityCreatorContext context = new EntityCreatorContext();
+        context.putUserData(EntityCreatorContext.DTO_TYPE, ENTITY_DTO_TYPE);
+        context.putUserData(EntityCreatorContext.ENTITY_ID, ENTITY_ID);
+        GenerationContextRegistry.getInstance().setContext(context);
+    }
+
     /**
      * Test generation of DeleteByIdCommand model for entity.
      */
diff --git a/tests/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityGeneratorTest.java b/tests/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityGeneratorTest.java
index 0554c016a..3150828c2 100644
--- a/tests/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityGeneratorTest.java
+++ b/tests/com/magento/idea/magento2plugin/actions/generation/generator/DeleteEntityGeneratorTest.java
@@ -6,18 +6,30 @@
 package com.magento.idea.magento2plugin.actions.generation.generator;
 
 import com.intellij.psi.PsiFile;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.DeleteEntityControllerFileData;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 import com.magento.idea.magento2plugin.magento.files.actions.DeleteActionFile;
 
 public class DeleteEntityGeneratorTest extends BaseGeneratorTestCase {
 
     private static final String MODULE_NAME = "Foo_Bar";
     private static final String ENTITY_NAME = "Company";
+    private static final String ENTITY_DTO_TYPE = "Foo\\Bar\\Model\\Data\\CompanyData";
     private static final String ACL = "Foo_Bar::company_id";
     private static final String ENTITY_ID = "company_id";
     private static final String EXPECTED_DIRECTORY =
             "/src/app/code/Foo/Bar/Controller/Adminhtml/" + ENTITY_NAME;
 
+    @Override
+    protected void setUp() throws Exception {
+        super.setUp();
+        final EntityCreatorContext context = new EntityCreatorContext();
+        context.putUserData(EntityCreatorContext.DTO_TYPE, ENTITY_DTO_TYPE);
+        context.putUserData(EntityCreatorContext.ENTITY_ID, ENTITY_ID);
+        GenerationContextRegistry.getInstance().setContext(context);
+    }
+
     /**
      * Test generation of Delete controller.
      */
diff --git a/tests/com/magento/idea/magento2plugin/actions/generation/generator/FormButtonBlockGeneratorTest.java b/tests/com/magento/idea/magento2plugin/actions/generation/generator/FormButtonBlockGeneratorTest.java
index 7a94f1ded..9445ef239 100644
--- a/tests/com/magento/idea/magento2plugin/actions/generation/generator/FormButtonBlockGeneratorTest.java
+++ b/tests/com/magento/idea/magento2plugin/actions/generation/generator/FormButtonBlockGeneratorTest.java
@@ -7,8 +7,10 @@
 
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiFile;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.UiComponentFormButtonData;
 import com.magento.idea.magento2plugin.actions.generation.generator.util.NamespaceBuilder;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 
 public class FormButtonBlockGeneratorTest extends BaseGeneratorTestCase {
 
@@ -17,10 +19,20 @@ public class FormButtonBlockGeneratorTest extends BaseGeneratorTestCase {
     private static final String ENTITY_NAME = "Book";
     private static final String BLOCK_DIRECTORY = "Block/Form/" + ENTITY_NAME;
     private static final String ACTION_NAME = "test";
+    private static final String ENTITY_DTO_TYPE = "Foo\\Bar\\Model\\Data\\BookData";
     private static final String ENTITY_ID = "book_id";
     public static final String EXPECTED_DIRECTORY
             = "src/app/code/Foo/Bar/Block/Form/" + ENTITY_NAME;
 
+    @Override
+    protected void setUp() throws Exception {
+        super.setUp();
+        final EntityCreatorContext context = new EntityCreatorContext();
+        context.putUserData(EntityCreatorContext.DTO_TYPE, ENTITY_DTO_TYPE);
+        context.putUserData(EntityCreatorContext.ENTITY_ID, ENTITY_ID);
+        GenerationContextRegistry.getInstance().setContext(context);
+    }
+
     /**
      * Test creation Block type Save.
      */
diff --git a/tests/com/magento/idea/magento2plugin/actions/generation/generator/FormGenericButtonBlockGeneratorTest.java b/tests/com/magento/idea/magento2plugin/actions/generation/generator/FormGenericButtonBlockGeneratorTest.java
index 6df775add..bf56fe0b5 100644
--- a/tests/com/magento/idea/magento2plugin/actions/generation/generator/FormGenericButtonBlockGeneratorTest.java
+++ b/tests/com/magento/idea/magento2plugin/actions/generation/generator/FormGenericButtonBlockGeneratorTest.java
@@ -5,7 +5,9 @@
 
 package com.magento.idea.magento2plugin.actions.generation.generator;
 
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.FormGenericButtonBlockData;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 import com.magento.idea.magento2plugin.magento.files.FormGenericButtonBlockFile;
 
 public class FormGenericButtonBlockGeneratorTest extends BaseGeneratorTestCase {
@@ -13,9 +15,19 @@ public class FormGenericButtonBlockGeneratorTest extends BaseGeneratorTestCase {
     private static final String MODULE_NAME = "Foo_Bar";
     private static final String ENTITY_NAME = "Book";
     private static final String ENTITY_ID = "book_id";
+    private static final String ENTITY_DTO_TYPE = "Foo\\Bar\\Model\\Data\\BookData";
     private static final String EXPECTED_DIRECTORY
             = "src/app/code/Foo/Bar/Block/Form/" + ENTITY_NAME;
 
+    @Override
+    protected void setUp() throws Exception {
+        super.setUp();
+        final EntityCreatorContext context = new EntityCreatorContext();
+        context.putUserData(EntityCreatorContext.DTO_TYPE, ENTITY_DTO_TYPE);
+        context.putUserData(EntityCreatorContext.ENTITY_ID, ENTITY_ID);
+        GenerationContextRegistry.getInstance().setContext(context);
+    }
+
     /**
      * Test generation of form ui component generic button block.
      */
diff --git a/tests/com/magento/idea/magento2plugin/actions/generation/generator/GridActionColumnFileGeneratorTest.java b/tests/com/magento/idea/magento2plugin/actions/generation/generator/GridActionColumnFileGeneratorTest.java
index e5943a706..7e1b780c7 100644
--- a/tests/com/magento/idea/magento2plugin/actions/generation/generator/GridActionColumnFileGeneratorTest.java
+++ b/tests/com/magento/idea/magento2plugin/actions/generation/generator/GridActionColumnFileGeneratorTest.java
@@ -5,7 +5,9 @@
 
 package com.magento.idea.magento2plugin.actions.generation.generator;
 
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.GridActionColumnData;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 import com.magento.idea.magento2plugin.magento.files.GridActionColumnFile;
 
 public class GridActionColumnFileGeneratorTest extends BaseGeneratorTestCase {
@@ -15,6 +17,16 @@ public class GridActionColumnFileGeneratorTest extends BaseGeneratorTestCase {
     private static final String ENTITY_ID_COLUMN = "book_id";
     private static final String EDIT_URL_PATH = "book_book_edit";
     private static final String DELETE_URL_PATH = "book_book_delete";
+    private static final String ENTITY_DTO_TYPE = "Foo\\Bar\\Model\\Data\\BookData";
+
+    @Override
+    protected void setUp() throws Exception {
+        super.setUp();
+        final EntityCreatorContext context = new EntityCreatorContext();
+        context.putUserData(EntityCreatorContext.DTO_TYPE, ENTITY_DTO_TYPE);
+        context.putUserData(EntityCreatorContext.ENTITY_ID, ENTITY_ID_COLUMN);
+        GenerationContextRegistry.getInstance().setContext(context);
+    }
 
     /**
      * Test generation of grid actions column file.
diff --git a/tests/com/magento/idea/magento2plugin/actions/generation/generator/ModuleResourceModelGeneratorTest.java b/tests/com/magento/idea/magento2plugin/actions/generation/generator/ModuleResourceModelGeneratorTest.java
index bbf36c78d..e0543fb29 100644
--- a/tests/com/magento/idea/magento2plugin/actions/generation/generator/ModuleResourceModelGeneratorTest.java
+++ b/tests/com/magento/idea/magento2plugin/actions/generation/generator/ModuleResourceModelGeneratorTest.java
@@ -7,22 +7,62 @@
 
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiFile;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.ResourceModelData;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 
 public class ModuleResourceModelGeneratorTest extends BaseGeneratorTestCase {
 
     private static final String EXPECTED_DIR = "src/app/code/Foo/Bar/Model/ResourceModel";
+    private static final String ENTITY_ID_COLUMN = "entity_id";
+    private static final String ENTITY_DTO_TYPE = "Foo\\Bar\\Model\\Data\\EntityData";
+
+    @Override
+    protected void setUp() throws Exception {
+        super.setUp();
+        final EntityCreatorContext context = new EntityCreatorContext();
+        context.putUserData(EntityCreatorContext.DTO_TYPE, ENTITY_DTO_TYPE);
+        context.putUserData(EntityCreatorContext.ENTITY_ID, ENTITY_ID_COLUMN);
+        GenerationContextRegistry.getInstance().setContext(context);
+    }
 
     /**
      * Test generation of resource model file.
      */
     public void testGenerateFile() {
+        GenerationContextRegistry.getInstance().setContext(null);
+        final Project project = myFixture.getProject();
+        final ResourceModelData resourceModelData = new ResourceModelData(
+                "Foo_Bar",
+                "my_table",
+                "TestResourceModel",
+                ENTITY_ID_COLUMN
+        );
+        final ModuleResourceModelGenerator generator = new ModuleResourceModelGenerator(
+                resourceModelData,
+                project
+        );
+        final PsiFile resourceModelFile = generator.generate("test");
+        final String filePath = this.getFixturePath("TestResourceModel.php");
+        final PsiFile expectedFile = myFixture.configureByFile(filePath);
+
+        assertGeneratedFileIsCorrect(
+                expectedFile,
+                EXPECTED_DIR,
+                resourceModelFile
+        );
+    }
+
+    /**
+     * Test generation of resource model file with reference to entity id column.
+     */
+    public void testGenerateFileWithDtoReference() {
         final Project project = myFixture.getProject();
         final ResourceModelData resourceModelData = new ResourceModelData(
                 "Foo_Bar",
                 "my_table",
                 "TestResourceModel",
-                "entity_id"
+                ENTITY_ID_COLUMN
         );
         final ModuleResourceModelGenerator generator = new ModuleResourceModelGenerator(
                 resourceModelData,
diff --git a/tests/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityActionGeneratorTest.java b/tests/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityActionGeneratorTest.java
index 2380e210b..9df4db4b0 100644
--- a/tests/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityActionGeneratorTest.java
+++ b/tests/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityActionGeneratorTest.java
@@ -18,7 +18,7 @@ public class SaveEntityActionGeneratorTest extends BaseGeneratorTestCase {
     private static final String EXPECTED_DIRECTORY =
             "/src/app/code/Foo/Bar/Controller/Adminhtml/" + ENTITY_NAME;
     private static final String ACL = "Foo_Bar::company_id";
-    private static final String ENTITY_ID = "entity_id";
+    private static final String ENTITY_ID = "company_id";
 
     /**
      * Test generation of Save controller for entity without interface.
diff --git a/tests/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityCommandGeneratorTest.java b/tests/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityCommandGeneratorTest.java
index 74607e64e..abc1c1eb2 100644
--- a/tests/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityCommandGeneratorTest.java
+++ b/tests/com/magento/idea/magento2plugin/actions/generation/generator/SaveEntityCommandGeneratorTest.java
@@ -12,6 +12,7 @@ public class SaveEntityCommandGeneratorTest extends BaseGeneratorTestCase {
 
     private static final String MODULE_NAME = "Foo_Bar";
     private static final String ENTITY_NAME = "Book";
+    private static final String ENTITY_ID = "book_id";
     private static final String MODEL_NAME = ENTITY_NAME + "Model";
     private static final String RESOURCE_MODEL_NAME = ENTITY_NAME + "Resource";
     private static final String DTO_NAME = ENTITY_NAME + "Data";
@@ -26,6 +27,7 @@ public void testGenerateSaveEntityCommandFile() {
         final SaveEntityCommandData saveEntityCommandData = new SaveEntityCommandData(
                 MODULE_NAME,
                 ENTITY_NAME,
+                ENTITY_ID,
                 MODEL_NAME,
                 RESOURCE_MODEL_NAME,
                 DTO_NAME,
diff --git a/tests/com/magento/idea/magento2plugin/actions/generation/generator/UiComponentGridDataProviderGeneratorTest.java b/tests/com/magento/idea/magento2plugin/actions/generation/generator/UiComponentGridDataProviderGeneratorTest.java
index 07d91ae3c..5b94a8cfd 100644
--- a/tests/com/magento/idea/magento2plugin/actions/generation/generator/UiComponentGridDataProviderGeneratorTest.java
+++ b/tests/com/magento/idea/magento2plugin/actions/generation/generator/UiComponentGridDataProviderGeneratorTest.java
@@ -7,21 +7,33 @@
 
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiFile;
+import com.magento.idea.magento2plugin.actions.generation.context.EntityCreatorContext;
 import com.magento.idea.magento2plugin.actions.generation.data.GetListQueryModelData;
 import com.magento.idea.magento2plugin.actions.generation.data.UiComponentDataProviderData;
+import com.magento.idea.magento2plugin.actions.generation.util.GenerationContextRegistry;
 
 public class UiComponentGridDataProviderGeneratorTest extends BaseGeneratorTestCase {
 
     private static final String EXPECTED_DIRECTORY = "src/app/code/Foo/Bar/Ui/Component/Listing";
     private static final String MODULE_NAME = "Foo_Bar";
     private static final String ENTITY_NAME = "Book";
-    private static final String ENTITY_ID_FIELD_NAME = "entity_id";
+    private static final String ENTITY_ID_FIELD_NAME = "book_id";
     private static final String MODEL_NAME = "Book";
     private static final String PROVIDER_CLASS_NAME = "GridDataProvider";
     private static final String PROVIDER_PATH = "Ui/Component/Listing";
+    private static final String ENTITY_DTO_TYPE = "Foo\\Bar\\Model\\Data\\BookData";
     private static final String COLLECTION_NAME = "Collection";
     private static final String ACL = "Foo_Bar::book_management";
 
+    @Override
+    protected void setUp() throws Exception {
+        super.setUp();
+        final EntityCreatorContext context = new EntityCreatorContext();
+        context.putUserData(EntityCreatorContext.DTO_TYPE, ENTITY_DTO_TYPE);
+        context.putUserData(EntityCreatorContext.ENTITY_ID, ENTITY_ID_FIELD_NAME);
+        GenerationContextRegistry.getInstance().setContext(context);
+    }
+
     /**
      * Test data provider class file generation with custom type.
      */