From 3c88d3679fd6e321b2505ed0d9eefde444b62adf Mon Sep 17 00:00:00 2001
From: Ilia Zakoulov <ilia.zakoulov@jetbrains.com>
Date: Mon, 12 May 2025 16:30:28 +0200
Subject: [PATCH] PY-77360: Introduce tests for PyDocStringGenerator

GitOrigin-RevId: f230b0f4069ea0c2c0032f4668051cfd23aa010a
---
 .../docGeneration/moduleLevel.after.py        | 13 +++++++
 python/testData/docGeneration/moduleLevel.py  |  9 +++++
 .../moduleLevelAfterComments.after.py         | 17 ++++++++
 .../docGeneration/moduleLevelAfterComments.py | 13 +++++++
 .../docstings/PyDocstringGeneratorTestCase.kt | 39 +++++++++++++++++++
 5 files changed, 91 insertions(+)
 create mode 100644 python/testData/docGeneration/moduleLevel.after.py
 create mode 100644 python/testData/docGeneration/moduleLevel.py
 create mode 100644 python/testData/docGeneration/moduleLevelAfterComments.after.py
 create mode 100644 python/testData/docGeneration/moduleLevelAfterComments.py
 create mode 100644 python/testSrc/com/jetbrains/python/documentation/docstings/PyDocstringGeneratorTestCase.kt

diff --git a/python/testData/docGeneration/moduleLevel.after.py b/python/testData/docGeneration/moduleLevel.after.py
new file mode 100644
index 000000000000..45211c514829
--- /dev/null
+++ b/python/testData/docGeneration/moduleLevel.after.py
@@ -0,0 +1,13 @@
+"""
+This is documentation.
+It has multiple lines.
+"""
+import os
+import sys
+
+
+class Test:
+  pass
+
+def foo():
+  pass
diff --git a/python/testData/docGeneration/moduleLevel.py b/python/testData/docGeneration/moduleLevel.py
new file mode 100644
index 000000000000..283be013523a
--- /dev/null
+++ b/python/testData/docGeneration/moduleLevel.py
@@ -0,0 +1,9 @@
+import os
+import sys
+
+
+class Test:
+  pass
+
+def foo():
+  pass
diff --git a/python/testData/docGeneration/moduleLevelAfterComments.after.py b/python/testData/docGeneration/moduleLevelAfterComments.after.py
new file mode 100644
index 000000000000..31e17dc74733
--- /dev/null
+++ b/python/testData/docGeneration/moduleLevelAfterComments.after.py
@@ -0,0 +1,17 @@
+#!/bin/sh
+
+# some comments
+"""
+This is documentation.
+It has multiple lines.
+"""
+
+import os
+import sys
+
+
+class Test:
+  pass
+
+def foo():
+  pass
diff --git a/python/testData/docGeneration/moduleLevelAfterComments.py b/python/testData/docGeneration/moduleLevelAfterComments.py
new file mode 100644
index 000000000000..08f88ca96128
--- /dev/null
+++ b/python/testData/docGeneration/moduleLevelAfterComments.py
@@ -0,0 +1,13 @@
+#!/bin/sh
+
+# some comments
+
+import os
+import sys
+
+
+class Test:
+  pass
+
+def foo():
+  pass
diff --git a/python/testSrc/com/jetbrains/python/documentation/docstings/PyDocstringGeneratorTestCase.kt b/python/testSrc/com/jetbrains/python/documentation/docstings/PyDocstringGeneratorTestCase.kt
new file mode 100644
index 000000000000..ed526da8ad9d
--- /dev/null
+++ b/python/testSrc/com/jetbrains/python/documentation/docstings/PyDocstringGeneratorTestCase.kt
@@ -0,0 +1,39 @@
+// 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.documentation.docstings
+
+import com.intellij.idea.TestFor
+import com.intellij.openapi.command.WriteCommandAction
+import com.intellij.testFramework.TestDataPath
+import com.jetbrains.python.documentation.docstrings.PyDocstringGenerator
+import com.jetbrains.python.fixtures.PyTestCase
+import com.jetbrains.python.psi.PyFile
+
+@TestFor(classes = [PyDocstringGenerator::class])
+@TestDataPath("\$CONTENT_ROOT/../testData/docGeneration")
+class PyDocstringGeneratorTestCase : PyTestCase() {
+  override fun getTestDataPath(): String = super.getTestDataPath() + "/docGeneration"
+
+  fun testModuleLevel() = doTestModuleLevelDocumentation()
+
+  fun testModuleLevelAfterComments() = doTestModuleLevelDocumentation()
+
+  private fun doTestModuleLevelDocumentation() {
+    val testName = getTestName(true)
+    val file = myFixture.configureByFile("$testName.py")
+    if (file !is PyFile) {
+      error("Incorrect file. Must be a Python file.")
+    }
+    WriteCommandAction.runWriteCommandAction(myFixture.project) {
+      val docstringGenerator = PyDocstringGenerator.forDocStringOwner(file)
+      docstringGenerator.buildAndInsert(DOCUMENTATION)
+    }
+    myFixture.checkResultByFile("$testName.after.py")
+  }
+
+  companion object {
+    private const val DOCUMENTATION = "\"\"\"\n" +
+                                      "This is documentation.\n" +
+                                      "It has multiple lines.\n" +
+                                      "\"\"\""
+  }
+}