diff --git a/courseSection/editingPsiElementsLesson/sortMethodsProgrammingTask/src/main/kotlin/org/jetbrains/academy/plugin/course/dev/ui/SortMethodsAction.kt b/courseSection/editingPsiElementsLesson/sortMethodsProgrammingTask/src/main/kotlin/org/jetbrains/academy/plugin/course/dev/ui/SortMethodsAction.kt
index 09ead12..ea54011 100644
--- a/courseSection/editingPsiElementsLesson/sortMethodsProgrammingTask/src/main/kotlin/org/jetbrains/academy/plugin/course/dev/ui/SortMethodsAction.kt
+++ b/courseSection/editingPsiElementsLesson/sortMethodsProgrammingTask/src/main/kotlin/org/jetbrains/academy/plugin/course/dev/ui/SortMethodsAction.kt
@@ -6,20 +6,20 @@ import com.intellij.openapi.actionSystem.CommonDataKeys
 import com.intellij.openapi.command.WriteCommandAction
 import org.jetbrains.academy.plugin.course.dev.access.sortMethods
 import org.jetbrains.kotlin.psi.KtClass
+import org.jetbrains.kotlin.psi.KtFile
 
 class SortMethodsAction : AnAction() {
     override fun actionPerformed(e: AnActionEvent) {
-        val project = e.project ?: return
-        val editor = e.getData(CommonDataKeys.EDITOR) ?: return
-        val psiFile = e.getData(CommonDataKeys.PSI_FILE) ?: return
+        val project = e.project
+        val editor = e.getData(CommonDataKeys.EDITOR)
+        val psiFile = e.getData(CommonDataKeys.PSI_FILE) as? KtFile
 
-        val caret = editor.caretModel.currentCaret
-        val element = psiFile.findElementAt(caret.offset) ?: return
+        if (project == null || editor == null || psiFile == null) return
 
-        val ktClass = element.parent as? KtClass ?: return
+        val caret = editor.caretModel.currentCaret
+        val element = psiFile.findElementAt(caret.offset)
 
-        // Call your method sorting function here
-        // sortMethodsInClass(ktClass)
+        val ktClass = element?.parent as? KtClass ?: return
 
         WriteCommandAction.runWriteCommandAction(project) {
             // Execute your sorting logic here
diff --git a/courseSection/editingPsiElementsLesson/sortMethodsProgrammingTask/test/org/jetbrains/academy/plugin/course/dev/access/Tests.kt b/courseSection/editingPsiElementsLesson/sortMethodsProgrammingTask/test/org/jetbrains/academy/plugin/course/dev/access/Tests.kt
index 63ac479..ada1050 100644
--- a/courseSection/editingPsiElementsLesson/sortMethodsProgrammingTask/test/org/jetbrains/academy/plugin/course/dev/access/Tests.kt
+++ b/courseSection/editingPsiElementsLesson/sortMethodsProgrammingTask/test/org/jetbrains/academy/plugin/course/dev/access/Tests.kt
@@ -26,8 +26,6 @@ class Test : BasePlatformTestCase() {
             
         """.trimIndent()
         val file = myFixture.configureByText("Person.kt", fileContent)
-        sortMethods(file)
-
         val classes = PsiTreeUtil.findChildrenOfType(file, KtClass::class.java)
         for (ktClass in classes) {
             val methods = ktClass.declarations.filterIsInstance<KtNamedFunction>()