PY-61304 Add searching in import elements in "conftest.py" file

In addition to searching for fixtures as functions in the "conftest.py" file, add also searching in import elements.


Merge-request: IJ-MR-109809
Merged-by: Egor Eliseev <Egor.Eliseev@jetbrains.com>

GitOrigin-RevId: fe2832b23db8b5fda461b1e07a129975b2dd57bf

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

@@ -61,20 +61,22 @@ data class NamedFixtureLink(val fixture: PyTestFixture, val importElement: PyImp
 private fun PyTestFixture.getContainingFile(): PsiFile? = function?.containingFile
 
 /**
- * Check if fixture is in the "conftest.py" file in the given directory
+ * Check if a fixture is in the "conftest.py" file in the given directory
  */
 private fun PyTestFixture.isInConftestInDir(directory: PsiDirectory): Boolean {
   getContainingFile()?.let { return it.containingDirectory == directory && it.name == CONFTEST_PY } ?: return false
 }
 
 /**
- * Searching the right fixture in
+ * Searching for the right fixture in
  * 1. [func] containing class and parent classes
  * 2. [func] containing file
- * 3. import statements
+ * 3. Import statements in fixture file
  * 4. "conftest.py" files in parent directories
+ * 5. Import statements in "conftest.py" file
+ * 6. Reserved fixtures in "_pytest" dir
  *
- * [fixtureCandidates] All pytest fixtures in project that could be used by [func].forWhat
+ * [fixtureCandidates] All pytest fixtures in a project that could be used by [func].forWhat
  * [func] PyFunction using [pyFixtureElement]
  * [pyFixtureElement] Fixture provided as PyNamedParameter or PyStringLiteralExpression
  * [projectPath] Project directory path
@@ -116,31 +118,19 @@ private fun findRightFixture(fixtureCandidates: List<PyTestFixture>,
 
   // search in import
   if (currentFile is PyFile) {
-    val importedFixture = currentFile.findExportedName(elementName) as? PyImportElement
-    val resolveImportElements = importedFixture?.multiResolve()?.map { it.element }
-    if (importedFixture != null) {
-      // if fixture is imported as `from module import some_fixture as sf`
-      resolveImportElements?.filterIsInstance<PyFunction>()?.firstOrNull()?.let { fixture ->
-        return NamedFixtureLink(PyTestFixture(func, fixture, fixture.name ?: ""), importedFixture)
-      }
-
-      resolveImportElements?.let { list ->
-        fixtureCandidates.find { fixture -> list.contains(fixture.function) }?.let {
-          return NamedFixtureLink(it, importedFixture)
-        }
-      }
-    }
+    getFixtureFromImports(currentFile, elementName, func, fixtureCandidates)?.let { return it }
   }
 
   // search in "conftest.py" in parents directories
   if (!fixtureCandidates.isEmpty()) {
     var currentDirectory = currentFile.containingDirectory
     while (currentDirectory != null && currentDirectory.virtualFile.path != projectPath) {
-      fixtureCandidates.find { it.isInConftestInDir(currentDirectory) }?.let { return NamedFixtureLink(it, null) }
+      searchInConftest(fixtureCandidates, currentDirectory, elementName, func)?.let { return it }
+
       currentDirectory.parentDirectory?.let { currentDirectory = it }
     }
     currentDirectory?.let {
-      fixtureCandidates.find { it.isInConftestInDir(currentDirectory) }?.let { return NamedFixtureLink(it, null) }
+      searchInConftest(fixtureCandidates, currentDirectory, elementName, func)?.let { return it }
     }
   }
 
@@ -158,6 +148,39 @@ private fun findRightFixture(fixtureCandidates: List<PyTestFixture>,
   return null
 }
 
+/**
+ * Search fixture or imported fixture in 'constest.py' file
+ */
+private fun searchInConftest(fixtureCandidates: List<PyTestFixture>, currentDirectory: PsiDirectory, elementName: String, func: PyFunction?): NamedFixtureLink? {
+  fixtureCandidates.find { it.isInConftestInDir(currentDirectory) }?.let { return NamedFixtureLink(it, null) }
+  // search in imports in "conftest.py" file
+  (currentDirectory.findFile(CONFTEST_PY) as? PyFile)?.let { pyFile ->
+    getFixtureFromImports(pyFile, elementName, func, fixtureCandidates)?.let { return it }
+  }
+  return null
+}
+
+/**
+ * Return fixture from import element
+ */
+private fun getFixtureFromImports(targetFile: PyFile, elementName: String, func: PyFunction?, fixtureCandidates: List<PyTestFixture>): NamedFixtureLink? {
+  val importedFixture = targetFile.findExportedName(elementName) as? PyImportElement
+  val resolveImportElements = importedFixture?.multiResolve()?.map { it.element }
+  if (importedFixture != null) {
+    // if fixture is imported as `from module import some_fixture as sf`
+    resolveImportElements?.filterIsInstance<PyFunction>()?.firstOrNull()?.let { fixture ->
+      return NamedFixtureLink(PyTestFixture(func, fixture, fixture.name ?: ""), importedFixture)
+    }
+
+    resolveImportElements?.let { list ->
+      fixtureCandidates.find { fixture -> list.contains(fixture.function) }?.let {
+        return NamedFixtureLink(it, importedFixture)
+      }
+    }
+  }
+  return null
+}
+
 /**
  * @return Boolean If named parameter has fixture or not
  */

python/testData/testPytestFixtureResolving/testImportedFixturesInConftest/conftest.py

@@ -0,0 +1,3 @@
+from fixtures.foo_fixtures import foo_fixture1, foo_fixture2
+
+__all__ = ['foo_fixture1', 'foo_fixture2']

python/testData/testPytestFixtureResolving/testImportedFixturesInConftest/fixtures/foo_fixtures.py

@@ -0,0 +1,11 @@
+import pytest
+
+
+@pytest.fixture()
+def foo_fixture1():
+    return 1
+
+
+@pytest.fixture()
+def foo_fixture2():
+    return 2

python/testData/testPytestFixtureResolving/testImportedFixturesInConftest/test_imported_fixtures_in_conftest.py

@@ -0,0 +1,6 @@
+import pytest
+
+
+def test_reproduce(foo_<caret>fixture1, foo_fixture2):
+    assert foo_fixture1 == 1
+    assert foo_fixture2 == 2

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

@@ -52,6 +52,12 @@ class PyTestFixtureResolvingTest : PyTestCase() {
     const val TEST_USE_FIXTURES_DECORATED_CLASS = "/test_decorator_for_class.py"
     const val TEST_USE_FIXTURES_DECORATED_TEST = "/test_decorator_for_test.py"
     const val TEST_USE_FIXTURES_EXPRESSION = "/test_expression.py"
+
+    const val IMPORTED_FIXTURE_IN_CONFTEST_DIR_NAME = "testImportedFixturesInConftest"
+    const val IMPORTED_FIXTURE_IN_CONFTEST_DIR = "/$IMPORTED_FIXTURE_IN_CONFTEST_DIR_NAME"
+    const val TEST_IMPORTED_FIXTURE_IN_CONFTEST = "/test_imported_fixtures_in_conftest.py"
+    const val IMPORTED_FIXTURE_IN_CONFTEST_FOO_FIXTURES_DIR_NAME = "fixtures"
+    const val IMPORTED_FIXTURE_IN_CONFTEST_FOO_FIXTURES = "foo_fixtures.py"
   }
 
   override fun getTestDataPath() = super.getTestDataPath() + TESTS_SUBDIR
@@ -213,4 +219,8 @@ class PyTestFixtureResolvingTest : PyTestCase() {
   fun testUseFixturesExpression() {
     assertCorrectFile(USE_FIXTURES_TEST_DIR, TEST_USE_FIXTURES_EXPRESSION, CONFTEST_PY, USE_FIXTURES_TEST_DIR_NAME)
   }
+
+  fun testImportedFixtureInConftest() {
+    assertCorrectFile(IMPORTED_FIXTURE_IN_CONFTEST_DIR, TEST_IMPORTED_FIXTURE_IN_CONFTEST, IMPORTED_FIXTURE_IN_CONFTEST_FOO_FIXTURES, IMPORTED_FIXTURE_IN_CONFTEST_FOO_FIXTURES_DIR_NAME)
+  }
 }