diff --git a/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixture.kt b/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixture.kt index bcdba98469f9..b3ad8eb3762e 100644 --- a/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixture.kt +++ b/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixture.kt @@ -8,10 +8,7 @@ import com.intellij.psi.stubs.StubIndex import com.intellij.psi.util.PsiTreeUtil import com.intellij.util.ThreeState import com.jetbrains.python.PyNames -import com.jetbrains.python.psi.PyDecorator -import com.jetbrains.python.psi.PyElement -import com.jetbrains.python.psi.PyFunction -import com.jetbrains.python.psi.PyNamedParameter +import com.jetbrains.python.psi.* import com.jetbrains.python.psi.impl.PyEvaluator import com.jetbrains.python.psi.stubs.PyDecoratorStubIndex import com.jetbrains.python.psi.types.TypeEvalContext @@ -19,7 +16,9 @@ import com.jetbrains.python.testing.PyTestFrameworkService import com.jetbrains.python.testing.TestRunnerService import com.jetbrains.python.testing.isTestElement -private const val decoratorName: String = "pytest.fixture" +private val decoratorNames = arrayOf("pytest.fixture", "fixture") + +private fun PyDecoratorList.hasDecorator(vararg names: String) = names.map { findDecorator(it) }.firstOrNull() != null ?: false /** * If named parameter has fixture -- return it @@ -35,10 +34,11 @@ internal fun getFixture(element: PyNamedParameter, typeEvalContext: TypeEvalCont */ fun PyNamedParameter.isFixture(typeEvalContext: TypeEvalContext) = getFixture(this, typeEvalContext) != null + /** * @return Boolean is function decorated as fixture or marked so by EP */ -internal fun PyFunction.isFixture() = decoratorList?.findDecorator(decoratorName) != null || isCustomFixture() +internal fun PyFunction.isFixture() = decoratorList?.hasDecorator(*decoratorNames) ?: false || isCustomFixture() /** @@ -49,11 +49,13 @@ internal fun PyFunction.isFixture() = decoratorList?.findDecorator(decoratorName data class PyTestFixture(val function: PyFunction? = null, val resolveTarget: PyElement? = function, val name: String) -fun findDecoratorsByName(module: Module, name: String): Iterable = - StubIndex.getElements(PyDecoratorStubIndex.KEY, name, module.project, - GlobalSearchScope.union( - arrayOf(module.moduleContentScope, GlobalSearchScope.moduleRuntimeScope(module, true))), - PyDecorator::class.java) +fun findDecoratorsByName(module: Module, vararg names: String): Iterable = + names.map { name -> + StubIndex.getElements(PyDecoratorStubIndex.KEY, name, module.project, + GlobalSearchScope.union( + arrayOf(module.moduleContentScope, GlobalSearchScope.moduleRuntimeScope(module, true))), + PyDecorator::class.java) + }.flatten() private fun createFixture(decorator: PyDecorator): PyTestFixture? { @@ -88,7 +90,7 @@ internal fun getFixtures(module: Module, forWhat: PyFunction, typeEvalContext: T forWhat.isSubjectForFixture() ) { //Fixtures - (findDecoratorsByName(module, decoratorName) + (findDecoratorsByName(module, *decoratorNames) .mapNotNull { createFixture(it) } + getCustomFixtures(typeEvalContext, forWhat)) .filterNot { fixture && it.name == forWhat.name }.toList() // Do not suggest fixture for itself