diff --git a/jadx-core/src/test/java/jadx/tests/api/IntegrationTest.java b/jadx-core/src/test/java/jadx/tests/api/IntegrationTest.java
index 44351b16dcd3cf817adfc05b679876518f590a6f..539e2120c415e5d2cdd15e329542b62e041b0cf0 100644
--- a/jadx-core/src/test/java/jadx/tests/api/IntegrationTest.java
+++ b/jadx-core/src/test/java/jadx/tests/api/IntegrationTest.java
@@ -21,6 +21,7 @@ import java.util.concurrent.TimeoutException;
 import java.util.jar.JarOutputStream;
 
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 
@@ -157,17 +158,23 @@ public abstract class IntegrationTest extends TestUtils {
 		assertThat("Class not found: " + clsName, cls, notNullValue());
 		assertThat(clsName, is(cls.getClassInfo().getFullName()));
 
-		decompileAndCheck(jadxDecompiler, Collections.singletonList(cls));
+		decompileAndCheck(cls);
 		return cls;
 	}
 
-	public ClassNode searchCls(List<ClassNode> list, String fullClsName) {
+	@Nullable
+	public ClassNode searchCls(List<ClassNode> list, String clsName) {
 		for (ClassNode cls : list) {
-			if (cls.getClassInfo().getFullName().equals(fullClsName)) {
+			if (cls.getClassInfo().getFullName().equals(clsName)) {
 				return cls;
 			}
 		}
-		fail("Class not found by name " + fullClsName + " in list: " + list);
+		for (ClassNode cls : list) {
+			if (cls.getClassInfo().getShortName().equals(clsName)) {
+				return cls;
+			}
+		}
+		fail("Class not found by name " + clsName + " in list: " + list);
 		return null;
 	}
 
@@ -187,7 +194,11 @@ public abstract class IntegrationTest extends TestUtils {
 		return d;
 	}
 
-	protected void decompileAndCheck(JadxDecompiler d, List<ClassNode> clsList) {
+	protected void decompileAndCheck(ClassNode cls) {
+		decompileAndCheck(Collections.singletonList(cls));
+	}
+
+	protected void decompileAndCheck(List<ClassNode> clsList) {
 		if (!unloadCls) {
 			clsList.forEach(cls -> cls.add(AFlag.DONT_UNLOAD_CLASS));
 		}
@@ -207,6 +218,14 @@ public abstract class IntegrationTest extends TestUtils {
 			clsList.forEach(this::printSmali);
 		}
 
+		runChecks(clsList);
+	}
+
+	public void runChecks(ClassNode cls) {
+		runChecks(Collections.singletonList(cls));
+	}
+
+	protected void runChecks(List<ClassNode> clsList) {
 		clsList.forEach(this::checkCode);
 		compile(clsList);
 		clsList.forEach(this::runAutoCheck);
diff --git a/jadx-core/src/test/java/jadx/tests/api/SmaliTest.java b/jadx-core/src/test/java/jadx/tests/api/SmaliTest.java
index a3693e9142f5c2e25fc3b9ca9eb13a503a3df4db..78a9c3cd55bbc5a754c7c05255ec1dc6e4a27dae 100644
--- a/jadx-core/src/test/java/jadx/tests/api/SmaliTest.java
+++ b/jadx-core/src/test/java/jadx/tests/api/SmaliTest.java
@@ -57,7 +57,7 @@ public abstract class SmaliTest extends IntegrationTest {
 		jadxDecompiler = loadFiles(collectSmaliFiles(getTestPkg(), getTestName()));
 		RootNode root = JadxInternalAccess.getRoot(jadxDecompiler);
 		List<ClassNode> classes = root.getClasses(false);
-		decompileAndCheck(jadxDecompiler, classes);
+		decompileAndCheck(classes);
 		return classes;
 	}
 
diff --git a/jadx-core/src/test/java/jadx/tests/api/utils/assertj/JadxClassNodeAssertions.java b/jadx-core/src/test/java/jadx/tests/api/utils/assertj/JadxClassNodeAssertions.java
index fc4e652c3e3bb92d8ff518d88913d87287ee1d95..4b14a2af6ff6d1dff8867c7b1dc1cd8f110ec03d 100644
--- a/jadx-core/src/test/java/jadx/tests/api/utils/assertj/JadxClassNodeAssertions.java
+++ b/jadx-core/src/test/java/jadx/tests/api/utils/assertj/JadxClassNodeAssertions.java
@@ -1,13 +1,14 @@
 package jadx.tests.api.utils.assertj;
 
-import org.assertj.core.api.AbstractAssert;
+import org.assertj.core.api.AbstractObjectAssert;
 
 import jadx.api.ICodeInfo;
 import jadx.core.dex.nodes.ClassNode;
+import jadx.tests.api.IntegrationTest;
 
 import static jadx.tests.api.utils.assertj.JadxAssertions.assertThat;
 
-public class JadxClassNodeAssertions extends AbstractAssert<JadxClassNodeAssertions, ClassNode> {
+public class JadxClassNodeAssertions extends AbstractObjectAssert<JadxClassNodeAssertions, ClassNode> {
 	public JadxClassNodeAssertions(ClassNode cls) {
 		super(cls, JadxClassNodeAssertions.class);
 	}
@@ -20,4 +21,17 @@ public class JadxClassNodeAssertions extends AbstractAssert<JadxClassNodeAsserti
 		assertThat(codeStr).isNotBlank();
 		return new JadxCodeAssertions(codeStr);
 	}
+
+	public JadxCodeAssertions reloadCode(IntegrationTest testInstance) {
+		isNotNull();
+		ICodeInfo code = actual.reloadCode();
+		assertThat(code).isNotNull();
+		String codeStr = code.getCodeStr();
+		assertThat(codeStr).isNotBlank();
+
+		JadxCodeAssertions codeAssertions = new JadxCodeAssertions(codeStr);
+		codeAssertions.print();
+		testInstance.runChecks(actual);
+		return codeAssertions;
+	}
 }
diff --git a/jadx-core/src/test/java/jadx/tests/integration/rename/TestAnonymousInline.java b/jadx-core/src/test/java/jadx/tests/integration/rename/TestAnonymousInline.java
index f819046bd0438ceddb7b8813ca641e7c0e84b30f..9f25dd5c8626c286b94089a5975809101871737c 100644
--- a/jadx-core/src/test/java/jadx/tests/integration/rename/TestAnonymousInline.java
+++ b/jadx-core/src/test/java/jadx/tests/integration/rename/TestAnonymousInline.java
@@ -23,12 +23,11 @@ public class TestAnonymousInline extends IntegrationTest {
 	@Test
 	public void test() {
 		ClassNode cls = getClassNode(TestCls.class);
-		assertThat(cls.getCode())
+		assertThat(cls).code()
 				.containsOnlyOnce("return new Runnable() {");
 
-		assertThat(cls.reloadCode())
+		assertThat(cls).reloadCode(this)
 				.removeBlockComments() // remove comment about inlined class
-				.print()
 				.containsOnlyOnce("return new Runnable() {")
 				.doesNotContain("AnonymousClass1");
 	}
diff --git a/jadx-core/src/test/java/jadx/tests/integration/rename/TestConstReplace.java b/jadx-core/src/test/java/jadx/tests/integration/rename/TestConstReplace.java
index 378fc28c07e7963be8d439e258953c0dccca032a..43b2724d0f2ad6909746c13a12d34d215becd306 100644
--- a/jadx-core/src/test/java/jadx/tests/integration/rename/TestConstReplace.java
+++ b/jadx-core/src/test/java/jadx/tests/integration/rename/TestConstReplace.java
@@ -20,11 +20,10 @@ public class TestConstReplace extends IntegrationTest {
 	@Test
 	public void test() {
 		ClassNode cls = getClassNode(TestCls.class);
-		assertThat(cls.getCode())
+		assertThat(cls).code()
 				.containsOnlyOnce("return CONST;");
 
-		assertThat(cls.reloadCode())
-				.print()
+		assertThat(cls).reloadCode(this)
 				.containsOnlyOnce("return CONST;");
 	}
 }
diff --git a/jadx-core/src/test/java/jadx/tests/integration/rename/TestFieldWithGenericRename.java b/jadx-core/src/test/java/jadx/tests/integration/rename/TestFieldWithGenericRename.java
index 25823dc10fc67df04e9e4f48cfb46482fc2300db..1852c41a5f49355e19f4689f4cfeda85eb602de8 100644
--- a/jadx-core/src/test/java/jadx/tests/integration/rename/TestFieldWithGenericRename.java
+++ b/jadx-core/src/test/java/jadx/tests/integration/rename/TestFieldWithGenericRename.java
@@ -18,10 +18,12 @@ public class TestFieldWithGenericRename extends IntegrationTest {
 	@Test
 	public void test() {
 		ClassNode cls = getClassNode(TestCls.class);
-		assertThat(cls.getCode()).containsOnlyOnce("List<String> list;");
+		assertThat(cls).code()
+				.containsOnlyOnce("List<String> list;");
 
 		cls.searchFieldByName("list").getFieldInfo().setAlias("listFieldRenamed");
 
-		assertThat(cls.reloadCode()).print().containsOnlyOnce("List<String> listFieldRenamed;");
+		assertThat(cls).reloadCode(this)
+				.containsOnlyOnce("List<String> listFieldRenamed;");
 	}
 }
diff --git a/jadx-core/src/test/java/jadx/tests/integration/rename/TestRenameEnum.java b/jadx-core/src/test/java/jadx/tests/integration/rename/TestRenameEnum.java
index 2ae89b18aa87cfd0810a24a55a08c686187501b6..51a6206d45cc9b2bfe2aa395bc9fb9e1536f41b5 100644
--- a/jadx-core/src/test/java/jadx/tests/integration/rename/TestRenameEnum.java
+++ b/jadx-core/src/test/java/jadx/tests/integration/rename/TestRenameEnum.java
@@ -30,14 +30,13 @@ public class TestRenameEnum extends IntegrationTest {
 	@Test
 	public void test() {
 		ClassNode cls = getClassNode(TestCls.class);
-		assertThat(cls.getCode())
+		assertThat(cls).code()
 				.containsOnlyOnce("public enum A ")
 				.containsOnlyOnce("ONE {");
 
 		cls.getInnerClasses().get(0).getClassInfo().changeShortName("ARenamed");
 
-		assertThat(cls.reloadCode())
-				.print()
+		assertThat(cls).reloadCode(this)
 				.containsOnlyOnce("public enum ARenamed ")
 				.containsOnlyOnce("ONE {");
 	}