PY-83395: Test add source root quick fix

GitOrigin-RevId: df736f6a7d5fb838d5947e2e43ee6a4c96626858

python/python-psi-impl/src/com/jetbrains/python/inspections/unresolvedReference/PyUnresolvedReferencesVisitor.java

@@ -365,7 +365,7 @@ public abstract class PyUnresolvedReferencesVisitor extends PyInspectionVisitor
       isAddedToInstallAllFix = true;
     }
 
-    ContainerUtil.addAll(fixes, getAddSourceRootQuickFix(node));
+    ContainerUtil.addIfNotNull(fixes, getAddSourceRootQuickFix(node));
 
     if (reference instanceof PySubstitutionChunkReference) {
       return;

python/testData/quickFixes/PyMarkDirectoryAsSourceRootQuickFixTest/mysrc/foo/abc.py

@@ -0,0 +1 @@
+from <error descr="Unresolved reference 'qwerty'">qwerty</error>.qwe import <error descr="Unresolved reference 'a_function'">a_function</error>

python/testData/quickFixes/PyMarkDirectoryAsSourceRootQuickFixTest/mysrc/foo/abc_after.py

@@ -0,0 +1 @@
+from qwerty.qwe import a_function

python/testData/quickFixes/PyMarkDirectoryAsSourceRootQuickFixTest/mysrc/foo/abc_no_src.py

@@ -0,0 +1 @@
+from <error descr="Unresolved reference 'qwejfknqweklfnjqlkjwnef'">qwejfknqweklfnjqlkjwnef</error>.qwe import <error descr="Unresolved reference 'a_function'">a_function</error>

python/testData/quickFixes/PyMarkDirectoryAsSourceRootQuickFixTest/mysrc/foo/abc_no_src_after.py

@@ -0,0 +1 @@
+from <error descr="Unresolved reference 'qwejfknqweklfnjqlkjwnef'">qwejfknqweklfnjqlkjwnef</error>.qwe import <error descr="Unresolved reference 'a_function'">a_function</error>

python/testData/quickFixes/PyMarkDirectoryAsSourceRootQuickFixTest/mysrc/qwerty/qwe.py

@@ -0,0 +1,5 @@
+def a_function():
+    pass
+
+def abc():
+    pass

python/testSrc/com/jetbrains/python/quickFixes/PyMarkDirectoryAsSourceRootQuickFixTest.kt

@@ -0,0 +1,49 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.quickFixes
+
+import com.intellij.openapi.roots.ModuleRootManager
+import com.intellij.testFramework.TestDataPath
+import com.jetbrains.python.PyQuickFixTestCase
+import com.jetbrains.python.inspections.unresolvedReference.PyUnresolvedReferencesInspection
+
+@TestDataPath("\$CONTENT_ROOT/../testData/quickFixes/PyMarkDirectoryAsSourceRootQuickFixTest")
+class PyMarkDirectoryAsSourceRootQuickFixTest: PyQuickFixTestCase() {
+  @Throws(Exception::class)
+  override fun setUp() {
+    super.setUp()
+    myFixture.copyDirectoryToProject("", "")
+  }
+
+  fun testUpdatesSourceRoots() = doSourceRootTest(pyFilePath = "mysrc/foo/abc.py", expectedSourceRootPath = "mysrc")
+
+  fun testNoQuickFixBecauseSourceRootNotFound() = doSourceRootTest(pyFilePath = "mysrc/foo/abc_no_src.py", expectedSourceRootPath = null)
+
+  private fun doSourceRootTest(
+    pyFilePath: String,
+    expectedSourceRootPath: String?
+  ) {
+    myFixture.enableInspections(PyUnresolvedReferencesInspection::class.java)
+    myFixture.configureByFile(pyFilePath)
+    myFixture.checkHighlighting(true, false, false)
+    val intentionAction = myFixture.filterAvailableIntentions(QUICK_FIX_NAME_BEGINNING)
+    if (expectedSourceRootPath == null) {
+      assertEmpty(intentionAction)
+      return
+    }
+    assertSize(1, intentionAction)
+    myFixture.launchAction(intentionAction[0])
+    myFixture.checkResultByFile(pyFilePath.removeSuffix(".py") + "_after.py", true)
+
+    val moduleRootManager = ModuleRootManager.getInstance(myFixture.module)
+    val sourceRoot = moduleRootManager.sourceRoots.firstOrNull {
+      it.path == TEST_PROJECT_ROOT_PATH + expectedSourceRootPath
+    }
+    assertNotNull("Source root '${expectedSourceRootPath}' not found.", sourceRoot)
+  }
+
+  companion object {
+    private const val TEST_PROJECT_ROOT_PATH = "/src/"
+    // QFIX.add.source.root.for.unresolved.import.name
+    private const val QUICK_FIX_NAME_BEGINNING = "Mark '"
+  }
+}