[python] PY-54771 Suggest fixtures in @pytest.mark.usefixtures("")

Merge-request: IJ-MR-173095
Merged-by: Morgan Bartholomew <morgan.bartholomew@jetbrains.com>

(cherry picked from commit f4a8479b755d71cd2932176acce98f4d2f8aaba5)

IJ-MR-173095

GitOrigin-RevId: 382ee0e32c3a819d28fda7f5ef5db575517474ff

python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt

@@ -1,10 +1,12 @@
 // 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.testing.pyTestFixtures
 
+import com.intellij.openapi.module.ModuleUtilCore
 import com.intellij.openapi.util.Ref
 import com.intellij.openapi.util.TextRange
 import com.intellij.patterns.PlatformPatterns
 import com.intellij.psi.*
+import com.intellij.psi.util.findParentOfType
 import com.intellij.util.ArrayUtil
 import com.intellij.util.ProcessingContext
 import com.intellij.util.containers.ContainerUtil
@@ -14,7 +16,9 @@ import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.COROUTINE
 import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.GENERATOR
 import com.jetbrains.python.psi.*
 import com.jetbrains.python.psi.resolve.ImportedResolveResult
+import com.jetbrains.python.psi.resolve.QualifiedNameFinder
 import com.jetbrains.python.psi.types.*
+import org.jetbrains.annotations.ApiStatus
 
 class PyTestFixtureReference(pyElement: PsiElement, fixture: PyTestFixture, private val importElement: PyElement? = null, range: TextRange? = null) : BaseReference(pyElement, range), PsiPolyVariantReference {
   private val functionRef = fixture.function?.let { SmartPointerManager.createPointer(it) }
@@ -33,6 +37,52 @@ class PyTestFixtureReference(pyElement: PsiElement, fixture: PyTestFixture, priv
     return resultList.toArray(emptyArray())
   }
 
+  @ApiStatus.Internal
+  override fun getVariants(): Array<Any> {
+    // Provide completion variants for fixtures, especially inside pytest.mark.usefixtures strings
+    val element = myElement
+    val project = element.project
+    val context = TypeEvalContext.codeCompletion(project, element.containingFile)
+
+    // If we are in a parameter, variants are handled elsewhere; be conservative and return empty
+    if (element !is PyStringLiteralExpression) return emptyArray()
+
+    val module = ModuleUtilCore.findModuleForPsiElement(element) ?: return emptyArray()
+
+    // Try to find the function we are attached to (decorated function)
+    val func = element.findParentOfType<PyFunction>()
+
+    val fixtureNames: List<String> = if (func != null) {
+      // For function-level decorators
+      getFixtures(module, func, context)
+        .filterNot { fixture ->
+          // `usefixtures` is only useful if no access to the fixture object is required.
+          // Otherwise, fixtures are provided via function arguments.
+          // Fixtures from pytest usually need some interaction to be useful.
+          fixture.function?.isFromPytestPackage() ?: false
+        }
+        .map { it.name }
+    }
+    else {
+      // For class-level decorators and module-level function calls
+      findDecoratorsByName(module, TEST_FIXTURE_DECORATOR_NAMES)
+        .filterNot { dec ->
+          // `usefixtures` is only useful if no access to the fixture object is required.
+          // Otherwise, fixtures are provided via function arguments.
+          // Fixtures from pytest usually need some interaction to be useful.
+          dec.target?.isFromPytestPackage() == true
+        }
+        .mapNotNull { dec -> getTestFixtureName(dec) ?: dec.target?.name }
+    }
+
+    // Filter out builtin/reserved pytest fixtures and the special 'request' pseudo-fixture
+    val filtered = fixtureNames.filterNot { name ->
+      name in reservedFixturesSet || name in reservedFixtureClassSet || name == REQUEST_FIXTURE
+    }
+
+    return filtered.distinct().sorted().toTypedArray()
+  }
+
   override fun handleElementRename(newElementName: String): PsiElement {
     if (myElement is PyStringLiteralExpression) {
       return myElement.replace(PyElementGenerator.getInstance(myElement.project).createStringLiteralFromString(newElementName))
@@ -85,7 +135,8 @@ private abstract class PyTestReferenceProvider : PsiReferenceProvider() {
 private object PyTestReferenceAsParameterProvider : PyTestReferenceProvider() {
   override fun getReferencesByElement(element: PsiElement, context: ProcessingContext): Array<PsiReference> {
     val namedParam = element as? PyNamedParameter ?: return emptyArray()
-    val namedFixtureParameterLink = getFixtureLink(namedParam, TypeEvalContext.codeAnalysis(element.project, element.containingFile)) ?: return emptyArray()
+    val namedFixtureParameterLink = getFixtureLink(namedParam, TypeEvalContext.codeAnalysis(element.project, element.containingFile))
+                                    ?: return emptyArray()
     val annotationLength = namedParam.annotation?.textLength ?: 0
     return arrayOf(PyTestFixtureReference(namedParam, namedFixtureParameterLink.fixture, namedFixtureParameterLink.importElement, TextRange(0, element.textLength - annotationLength)))
   }
@@ -93,7 +144,7 @@ private object PyTestReferenceAsParameterProvider : PyTestReferenceProvider() {
 
 private object PyTestReferenceAsStringProvider : PyTestReferenceProvider() {
   override fun getReferencesByElement(element: PsiElement, context: ProcessingContext): Array<PsiReference> {
-    if (element !is PyExpression) return emptyArray()
+    if (element !is PyStringLiteralExpression) return emptyArray()
     val argumentList = element.parent as? PyArgumentList ?: return emptyArray()
     val callExpr = argumentList.parent as? PyCallExpression ?: return emptyArray()
     var shouldProvide = false
@@ -113,11 +164,23 @@ private object PyTestReferenceAsStringProvider : PyTestReferenceProvider() {
     // else
     if (!shouldProvide) return emptyArray()
 
-    val namedFixtureParameterLink = getFixtureLink(element, TypeEvalContext.codeAnalysis(element.project, element.containingFile)) ?: return emptyArray()
-    return arrayOf(PyTestFixtureReference(element, namedFixtureParameterLink.fixture, namedFixtureParameterLink.importElement, TextRange(1, element.textLength - 1)))
+    val namedFixtureParameterLink = getFixtureLink(element, TypeEvalContext.codeAnalysis(element.project, element.containingFile))
+
+    return arrayOf(
+      if (namedFixtureParameterLink != null)
+        PyTestFixtureReference(element, namedFixtureParameterLink.fixture, namedFixtureParameterLink.importElement, TextRange(1, element.textLength - 1))
+      else
+      // Provide a soft reference to enable completion even when nothing is resolvable yet (e.g., empty string)
+        PyTestFixtureReference(element, PyTestFixture(null, null, element.stringValue), null, TextRange(1, element.textLength - 1))
+    )
   }
 }
 
+private fun PyElement.isFromPytestPackage(): Boolean =
+  QualifiedNameFinder.findShortestImportableQName(containingFile)?.firstComponent.let {
+    it == "pytest" || it == "_pytest"
+  }
+
 class PyTestFixtureReferenceContributor : PsiReferenceContributor() {
 
   override fun registerReferenceProviders(registrar: PsiReferenceRegistrar) {

python/testData/testCompletion/test_usefixtures_completion.py

@@ -0,0 +1,11 @@
+import pytest
+
+
+@pytest.fixture
+def my_fixture():
+    return True
+
+
+@pytest.mark.usefixtures("<caret>")
+def test_abc():
+    ...

python/testData/testCompletion/test_usefixtures_completion_class.py

@@ -0,0 +1,12 @@
+import pytest
+
+
+@pytest.fixture
+def my_fixture():
+    return True
+
+
+@pytest.mark.usefixtures("<caret>")
+class TestSuite:
+    def test_abc(self):
+        ...

python/testData/testCompletion/test_usefixtures_completion_module.py

@@ -0,0 +1,13 @@
+import pytest
+
+
+@pytest.fixture
+def my_fixture():
+    return True
+
+
+# Module-level usefixtures via pytestmark
+pytestmark = pytest.mark.usefixtures("<caret>")
+
+def test_abc():
+    ...

python/testSrc/com/jetbrains/python/testing/PyTestFixtureAndParametrizedTest.kt

@@ -1,6 +1,7 @@
 // Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.jetbrains.python.testing
 
+import com.intellij.idea.TestFor
 import com.intellij.testFramework.fixtures.CodeInsightTestFixture
 import com.jetbrains.python.fixtures.PyTestCase
 import com.jetbrains.python.inspections.unusedLocal.PyUnusedLocalInspection
@@ -45,6 +46,42 @@ class PyTestFixtureAndParametrizedTest : PyTestCase() {
     myFixture.checkResultByFile("after_test_test.txt")
   }
 
+  @TestFor(issues = ["PY-54771"])
+  fun `test usefixtures completion suggests fix on function`() {
+    myFixture.copyDirectoryToProject(".", ".")
+    myFixture.configureByFile("test_usefixtures_completion.py")
+    myFixture.completeBasic()
+    val variants = myFixture.lookupElementStrings
+    assertNotNull(variants)
+    assertTrue(variants!!.contains("my_fixture"))
+    // Built-in pytest fixtures should not be suggested
+    assertFalse(variants.contains("tmp_path"))
+  }
+
+  @TestFor(issues = ["PY-54771"])
+  fun `test usefixtures completion suggests fix on class`() {
+    myFixture.copyDirectoryToProject(".", ".")
+    myFixture.configureByFile("test_usefixtures_completion_class.py")
+    myFixture.completeBasic()
+    val variants = myFixture.lookupElementStrings
+    assertNotNull(variants)
+    assertTrue(variants!!.contains("my_fixture"))
+    // Built-in pytest fixtures should not be suggested
+    assertFalse(variants.contains("tmp_path"))
+  }
+
+  @TestFor(issues = ["PY-54771"])
+  fun `test usefixtures completion suggests fix on module`() {
+    myFixture.copyDirectoryToProject(".", ".")
+    myFixture.configureByFile("test_usefixtures_completion_module.py")
+    myFixture.completeBasic()
+    val variants = myFixture.lookupElementStrings
+    assertNotNull(variants)
+    assertTrue(variants!!.contains("my_fixture"))
+    // Built-in pytest fixtures should not be suggested
+    assertFalse(variants.contains("tmp_path"))
+  }
+
   fun testRename() {
     myFixture.configureByFile("test_for_rename.py")
     myFixture.renameElementAtCaret("spam")