From 4dd725cf8dec211c64b326fa41dca556052709da Mon Sep 17 00:00:00 2001
From: Leutrim Shala <Leutrim.shala@aisec.fraunhofer.de>
Date: Wed, 25 Sep 2024 14:28:23 +0200
Subject: [PATCH] Refactor test

---
 .../statementHandler/StatementHandlerTest.kt  | 21 ++++++++++---------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/statementHandler/StatementHandlerTest.kt b/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/statementHandler/StatementHandlerTest.kt
index df9b2a1ab3..1023f22fc7 100644
--- a/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/statementHandler/StatementHandlerTest.kt
+++ b/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/statementHandler/StatementHandlerTest.kt
@@ -161,15 +161,16 @@ class StatementHandlerTest : BaseTest() {
     @Test
     fun testTypeHints() {
         analyzeFile("type_hints.py")
-
-        // type comments
-        val a = result.refs["a"]
-        assertNotNull(a)
-        assertEquals(with(result) { assertResolvedType("int") }, a.type)
-
-        // type annotation
-        val b = result.refs["b"]
-        assertNotNull(b)
-        assertEquals(with(result) { assertResolvedType("str") }, b.type)
+        with(result) {
+            // type comments
+            val a = result.refs["a"]
+            assertNotNull(a)
+            assertEquals(assertResolvedType("int"), a.type)
+
+            // type annotation
+            val b = result.refs["b"]
+            assertNotNull(b)
+            assertEquals(assertResolvedType("str"), b.type)
+        }
     }
 }