diff --git a/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/DelegatingTestClassModel.java b/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/DelegatingTestClassModel.java
index ddbf4ade530..92e5603e98e 100644
--- a/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/DelegatingTestClassModel.java
+++ b/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/DelegatingTestClassModel.java
@@ -67,10 +67,4 @@ public class DelegatingTestClassModel extends TestClassModel {
     public Collection<AnnotationModel> getAnnotations() {
         return delegate.getAnnotations();
     }
-
-    @NotNull
-    @Override
-    public Collection<Class<?>> getOwnImports() {
-        return delegate.getOwnImports();
-    }
 }
diff --git a/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/SimpleTestClassModel.java b/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/SimpleTestClassModel.java
index 88e929a5171..e65a65e827d 100644
--- a/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/SimpleTestClassModel.java
+++ b/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/SimpleTestClassModel.java
@@ -236,12 +236,6 @@ public class SimpleTestClassModel extends TestClassModel {
         return annotations;
     }
 
-    @NotNull
-    @Override
-    public Collection<Class<?>> getOwnImports() {
-        return annotations.stream().map(AnnotationModel::getAnnotation).collect(Collectors.toSet());
-    }
-
     private class TestAllFilesPresentMethodModel extends TestMethodModel {
         @NotNull
         @Override
diff --git a/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/SingleClassTestModel.java b/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/SingleClassTestModel.java
index 26ce32d118b..0cd91251b12 100644
--- a/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/SingleClassTestModel.java
+++ b/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/SingleClassTestModel.java
@@ -146,12 +146,6 @@ public class SingleClassTestModel extends TestClassModel {
         return annotations;
     }
 
-    @NotNull
-    @Override
-    public Collection<Class<?>> getOwnImports() {
-        return annotations.stream().map(AnnotationModel::getAnnotation).collect(Collectors.toSet());
-    }
-
     private class TestAllFilesPresentMethodModel extends TestMethodModel {
         @NotNull
         @Override
diff --git a/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/TestGenerator.kt b/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/TestGenerator.kt
index dba2acfb2ca..454847d72f8 100644
--- a/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/TestGenerator.kt
+++ b/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/TestGenerator.kt
@@ -120,10 +120,7 @@ class TestGenerator(
                 override val annotations: Collection<AnnotationModel>
                     get() = emptyList()
 
-                override val ownImports: Collection<Class<*>>
-                    get() = emptyList()
-
-                override val imports: Collection<Class<*>>
+                override val imports: Set<Class<*>>
                     get() = super.imports
             }
         }
diff --git a/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/TestModel.kt b/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/TestModel.kt
index 47cacb7a9a9..e91c9ddfb58 100644
--- a/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/TestModel.kt
+++ b/generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/TestModel.kt
@@ -24,20 +24,19 @@ interface TestEntityModel {
 }
 
 interface ClassModel : TestEntityModel {
-    val ownImports: Collection<Class<*>>
     val innerTestClasses: Collection<TestClassModel>
     val methods: Collection<MethodModel>
     val isEmpty: Boolean
     val dataPathRoot: String?
     val annotations: Collection<AnnotationModel>
-    val imports: Collection<Class<*>>
+    val imports: Set<Class<*>>
 }
 
 abstract class TestClassModel : ClassModel {
-    override val imports: Collection<Class<*>>
+    override val imports: Set<Class<*>>
         get() {
-            return mutableListOf<Class<*>>().also { allImports ->
-                allImports.addAll(ownImports)
+            return mutableSetOf<Class<*>>().also { allImports ->
+                annotations.mapTo(allImports) { it.annotation }
                 methods.flatMapTo(allImports) { it.imports() }
                 innerTestClasses.flatMapTo(allImports) { it.imports }
             }