PY-82575 [PyCharm] Thread leaked in PipPackageManagerTest

Signed-off-by: Nikita.Ashihmin <nikita.ashihmin@jetbrains.com>

GitOrigin-RevId: ec14d39a1818b606642d2a6babfad9a354b5b936

python/src/com/jetbrains/python/packaging/management/PythonPackageManager.kt

@@ -4,6 +4,7 @@
 package com.jetbrains.python.packaging.management
 
 import com.intellij.execution.ExecutionException
+import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.service
 import com.intellij.openapi.diagnostic.thisLogger
@@ -11,6 +12,7 @@ import com.intellij.openapi.progress.runBlockingMaybeCancellable
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.util.Key
+import com.intellij.util.cancelOnDispose
 import com.intellij.util.messages.Topic
 import com.jetbrains.python.errorProcessing.PyResult
 import com.jetbrains.python.getOrNull
@@ -38,11 +40,13 @@ import kotlin.coroutines.cancellation.CancellationException
  * @see com.jetbrains.python.packaging.management.ui.PythonPackageManagerUI to execute commands with UI handlers
  */
 @ApiStatus.Experimental
-abstract class PythonPackageManager(val project: Project, val sdk: Sdk) {
+abstract class PythonPackageManager(val project: Project, val sdk: Sdk) : Disposable.Default {
   private val isInited = AtomicBoolean(false)
   private val initializationJob by lazy {
     PyPackageCoroutine.launch(project, start = CoroutineStart.LAZY) {
       initManager()
+    }.also {
+      it.cancelOnDispose(this)
     }
   }
 
@@ -123,7 +127,7 @@ abstract class PythonPackageManager(val project: Project, val sdk: Sdk) {
       installedPackages = packages
       PyPackageCoroutine.launch(project) {
         reloadOutdatedPackages()
-      }
+      }.cancelOnDispose(this)
 
       ApplicationManager.getApplication().messageBus.apply {
         syncPublisher(PACKAGE_MANAGEMENT_TOPIC).packagesChanged(sdk)

python/src/com/jetbrains/python/packaging/management/PythonPackageManagerServiceImpl.kt

@@ -7,6 +7,7 @@ import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.util.Disposer
 import com.jetbrains.python.packaging.bridge.PythonPackageManagementServiceBridge
 import com.jetbrains.python.packaging.requirementsTxt.PythonRequirementTxtSdkUtils
+import com.jetbrains.python.packaging.utils.PyPackageCoroutine
 import com.jetbrains.python.sdk.PythonSdkAdditionalData
 import com.jetbrains.python.sdk.getOrCreateAdditionalData
 import kotlinx.coroutines.CoroutineScope
@@ -26,6 +27,7 @@ internal class PythonPackageManagerServiceImpl(private val serviceScope: Corouti
 
     return cache.computeIfAbsent(cacheKey) {
       val createdSdk = PythonPackageManagerProvider.EP_NAME.extensionList.firstNotNullOf { it.createPackageManagerForSdk(project, sdk) }
+      Disposer.register(PyPackageCoroutine.getInstance(project), createdSdk)
       PythonRequirementTxtSdkUtils.migrateRequirementsTxtPathFromModuleToSdk(project, sdk)
       createdSdk
     }

python/src/com/jetbrains/python/packaging/utils/PyPackageCoroutine.kt

@@ -13,6 +13,7 @@ import kotlin.coroutines.CoroutineContext
 @ApiStatus.Internal
 class PyPackageCoroutine(val project: Project, val coroutineScope: CoroutineScope) : Disposable.Default {
   companion object {
+    fun getInstance(project: Project): PyPackageCoroutine = project.service<PyPackageCoroutine>()
     fun launch(project: Project?, context: CoroutineContext = Dispatchers.Main, start: CoroutineStart = CoroutineStart.DEFAULT, body: suspend CoroutineScope.() -> Unit): Job {
       return project?.service<PyPackageCoroutine>()?.coroutineScope?.launch(context, block = body, start = start)
              ?: ApplicationManager.getApplication().service<PyAppCoroutine>().coroutineScope.launch(context, block = body, start = start)

python/testSrc/com/jetbrains/env/python/PipPackageManagerTestBase.kt

@@ -1,10 +1,12 @@
 // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.jetbrains.env.python
 
+import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.use
 import com.intellij.testFramework.ProjectRule
 import com.intellij.testFramework.common.timeoutRunBlocking
+import com.intellij.testFramework.junit5.fixture.disposableFixture
 import com.jetbrains.python.packaging.pip.PipPythonPackageManager
-import kotlinx.coroutines.test.runTest
 import org.hamcrest.MatcherAssert.assertThat
 import org.hamcrest.Matchers.empty
 import org.hamcrest.Matchers.not
@@ -24,9 +26,12 @@ abstract class PipPackageManagerTestBase {
 
   @Test
   fun testList(): Unit =  timeoutRunBlocking(5.minutes) {
-    PipPythonPackageManager(projectRule.project, sdkRule.sdk).apply {
-      assertThat("No packages return", reloadPackages().successOrNull, not(empty()))
-      assertThat("Installed packages shouldn't be empty", listInstalledPackages(), not(empty()))
+    disposableFixture().get().use {
+      PipPythonPackageManager(projectRule.project, sdkRule.sdk).apply {
+        Disposer.register(it, this)
+        assertThat("No packages return", reloadPackages().successOrNull, not(empty()))
+        assertThat("Installed packages shouldn't be empty", listInstalledPackages(), not(empty()))
+      }
     }
   }
 }

python/testSrc/com/jetbrains/python/packaging/management/TestPythonPackageManagerService.kt

@@ -3,6 +3,7 @@ package com.jetbrains.python.packaging.management
 
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.util.Disposer
 import com.intellij.testFramework.replaceService
 import com.jetbrains.python.packaging.bridge.PythonPackageManagementServiceBridge
 import com.jetbrains.python.packaging.common.PythonPackage
@@ -16,14 +17,16 @@ class TestPythonPackageManagerService(val installedPackages: List<PythonPackage>
 
   override fun forSdk(project: Project, sdk: Sdk): PythonPackageManager {
     installedPackages.ifEmpty {
-      return TestPythonPackageManager(project, sdk)
+      return TestPythonPackageManager(project, sdk).also { Disposer.register(project, it) }
     }
 
     return TestPythonPackageManager(project, sdk)
       .withPackageInstalled(installedPackages)
       .withPackageNames(installedPackages.map { it.name })
       .withPackageDetails(PythonSimplePackageDetails(installedPackages.first().name, listOf(installedPackages.first().version),
-                                                     TestPackageRepository(installedPackages.map { it.name }.toSet())))
+                                                     TestPackageRepository(installedPackages.map { it.name }.toSet()))).also {
+        Disposer.register(project, it)
+      }
   }
 
   override fun bridgeForSdk(project: Project, sdk: Sdk): PythonPackageManagementServiceBridge {