diff --git a/.circleci/config.yml b/.circleci/config.yml
index 25861837c..501dc67c4 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -209,7 +209,7 @@ commands:
 # See: https://circleci.com/docs/jobs-steps/#jobs-overview & https://circleci.com/docs/configuration-reference/#jobs
 jobs:
   unit-tests:
-    parallelism: 3
+    parallelism: 8
     executor: default_image
     resource_class: "large"
     steps:
diff --git a/pyproject.toml b/pyproject.toml
index d74463475..f24df39ba 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -200,7 +200,7 @@ log_cli_level = "INFO"
 xfail_strict = true
 junit_duration_report = "call"
 junit_logging = "all"
-
+tmp_path_retention_policy = "failed"
 [build-system]
 requires = ["hatchling>=1.26.3", "hatch-vcs>=0.4.0", "setuptools-scm>=8.0.0"]
 build-backend = "hatchling.build"
diff --git a/ruff.toml b/ruff.toml
index 64eb33912..338b939e4 100644
--- a/ruff.toml
+++ b/ruff.toml
@@ -6,7 +6,20 @@ exclude = [
 ]
 unsafe-fixes = true
 [lint]
-select = ["F", "E", "W", "I", "UP", "D", "RUF", "DTZ", "TC", "EM", "TRY400"]
+select = [
+  "F",
+  "E",
+  "W",
+  "I",
+  "UP",
+  "D",
+  "YTT",
+  "RUF",
+  "DTZ",
+  "TC",
+  "EM",
+  "TRY400",
+]
 ignore = [
   "D100",
   "D101",
@@ -216,3 +229,5 @@ extend-generics = [
 ]
 [lint.isort]
 known-first-party = ["codegen"]
+[format]
+docstring-code-format = true
diff --git a/scripts/profiling/profile.py b/scripts/profiling/profile.py
index 4b0264a3e..a32e0282a 100644
--- a/scripts/profiling/profile.py
+++ b/scripts/profiling/profile.py
@@ -12,7 +12,7 @@ def profile(repo: str, memory: bool = False):
     output = f"{base}/raw.austin"
     compressed = f"{base}/compressed.austin"
     image = f"{base}/parse.svg"
-    test = Popen(["pytest", "tests/codemod/test_parse.py", "--extra-repos=true", "--durations=100", "-k", repo])
+    test = Popen(["pytest", "tests/integration/codemod/test_parse.py", "--extra-repos=true", "--durations=100", "-k", repo])
     try:
         command = ["sudo", "austin", "-p", str(test.pid), "-o", output]
         if memory:
diff --git a/src/codegen/sdk/core/detached_symbols/function_call.py b/src/codegen/sdk/core/detached_symbols/function_call.py
index 7699343e9..cb8d1ba5d 100644
--- a/src/codegen/sdk/core/detached_symbols/function_call.py
+++ b/src/codegen/sdk/core/detached_symbols/function_call.py
@@ -42,7 +42,7 @@ class FunctionCall(Expression[Parent], HasName, Resolvable, Generic[Parent]):
     """Abstract representation of a function invocation, e.g. in Python:
     ```
     def f():
-        g() # FunctionCall
+        g()  # FunctionCall
     ```
     """
 
diff --git a/src/codegen/sdk/core/import_resolution.py b/src/codegen/sdk/core/import_resolution.py
index 432ec2ab5..1c35dcf2c 100644
--- a/src/codegen/sdk/core/import_resolution.py
+++ b/src/codegen/sdk/core/import_resolution.py
@@ -403,6 +403,7 @@ def is_dynamic(self) -> bool:
             def my_function():
                 import foo  # Dynamic - only imported when function runs
 
+
             if condition:
                 from bar import baz  # Dynamic - only imported if condition is True
 
diff --git a/src/codegen/sdk/extensions/utils.pyx b/src/codegen/sdk/extensions/utils.pyx
index d46a377bc..e95b69ef4 100644
--- a/src/codegen/sdk/extensions/utils.pyx
+++ b/src/codegen/sdk/extensions/utils.pyx
@@ -54,6 +54,7 @@ def find_all_descendants(node: TSNode, type_names: Iterable[str] | str, max_dept
 def iter_all_descendants(node: TSNode, type_names: Iterable[str] | str, max_depth: int | None = None, nested: bool = True) -> Generator[TSNode, None, None]:
     if isinstance(type_names, str):
         type_names = [type_names]
+    type_names = frozenset(type_names)
 
     def traverse(current_node: TSNode, depth=0):
         if max_depth is not None and depth > max_depth:
diff --git a/src/codegen/sdk/python/file.py b/src/codegen/sdk/python/file.py
index 4a38f4e06..36595e165 100644
--- a/src/codegen/sdk/python/file.py
+++ b/src/codegen/sdk/python/file.py
@@ -56,7 +56,7 @@ def symbol_can_be_added(self, symbol: PySymbol) -> bool:
     @noapidoc
     @commiter
     def _parse_imports(self) -> None:
-        for import_node in iter_all_descendants(self.ts_node, {"import_statement", "import_from_statement", "future_import_statement"}):
+        for import_node in iter_all_descendants(self.ts_node, frozenset({"import_statement", "import_from_statement", "future_import_statement"})):
             PyImportStatement(import_node, self.node_id, self.G, self.code_block, 0)
 
     ####################################################################################################################
diff --git a/tests/conftest.py b/tests/conftest.py
index df179d306..771e04a9e 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,4 +1,3 @@
-import logging
 import os
 from pathlib import Path
 
@@ -17,7 +16,7 @@ def find_dirs_to_ignore(start_dir, prefix):
     return dirs_to_ignore
 
 
-def pytest_addoption(parser):
+def pytest_addoption(parser) -> None:
     parser.addoption(
         "--size",
         action="append",
@@ -68,15 +67,15 @@ def pytest_addoption(parser):
 
 
 # content of conftest.py
-def pytest_configure(config):
-    worker_id = os.environ.get("PYTEST_XDIST_WORKER")
-    if worker_id is not None:
-        os.makedirs("build/logs", exist_ok=True)
-        logging.basicConfig(
-            format=config.getini("log_file_format"),
-            filename=f"build/logs/tests_{worker_id}.log",
-            level=config.getini("log_file_level"),
-        )
+# def pytest_configure(config) -> None:
+#     worker_id = os.environ.get("PYTEST_XDIST_WORKER")
+#     if worker_id is not None:
+#         os.makedirs("build/logs", exist_ok=True)
+#         logging.basicConfig(
+#             format=config.getini("log_file_format"),
+#             filename=f"build/logs/tests_{worker_id}.log",
+#             level=config.getini("log_file_level"),
+#         )
 
 
 def is_git_lfs_pointer(file_path: Path) -> bool:
@@ -100,18 +99,24 @@ def pytest_runtest_makereport(item, call):
             raise RuntimeError(msg)
 
 
-@pytest.fixture(autouse=True)
-def skip_lfs_tests(request):
-    """Skip tests that depend on git LFS files if they haven't been pulled"""
-    # Lets not run if we are in CI
-    if os.getenv("CI") == "true" or os.getenv("CIRCLECI") == "true":
-        return
+# Lets not run if we are in CI
+
+
+IS_CI = os.getenv("CI") == "true" or os.getenv("CIRCLECI") == "true"
 
+
+@pytest.fixture(autouse=IS_CI)
+def skip_lfs_tests(request) -> None:
+    """Skip tests that depend on git LFS files if they haven't been pulled"""
     # Get the test module path
     test_path = Path(request.module.__file__)
 
     # Only run for integration tests
-    if not str(test_path).startswith(str(Path.cwd() / "tests" / "integration")):
+    try:
+        cwd = Path.cwd()
+    except FileNotFoundError:
+        return
+    if not str(test_path).startswith(str(cwd / "tests" / "integration")):
         return
 
     try:
diff --git a/tests/unit/codegen/cli/conftest.py b/tests/integration/codegen/cli/conftest.py
similarity index 76%
rename from tests/unit/codegen/cli/conftest.py
rename to tests/integration/codegen/cli/conftest.py
index 3263638ca..021975826 100644
--- a/tests/unit/codegen/cli/conftest.py
+++ b/tests/integration/codegen/cli/conftest.py
@@ -2,11 +2,13 @@
 import shutil
 import subprocess
 from pathlib import Path
+from unittest.mock import patch
 
 import pytest
 from click.testing import CliRunner
 
 from codegen.cli.commands.init.main import init_command
+from codegen.cli.workspace.venv_manager import VenvManager
 
 
 @pytest.fixture
@@ -17,7 +19,11 @@ def sample_repository(tmp_path: Path):
     subprocess.run(["git", "config", "--local", "user.name", "Test User"], check=True)
     subprocess.run(["git", "commit", "--allow-empty", "-m", "Initial commit"], check=True)
     subprocess.run(["git", "remote", "add", "origin", "https://github.com/test/test.git"], check=True)
-    return tmp_path
+    yield tmp_path
+    try:
+        shutil.rmtree(tmp_path)
+    except FileNotFoundError:
+        pass
 
 
 @pytest.fixture()
@@ -28,8 +34,10 @@ def runner():
 @pytest.fixture
 def initialized_repo(sample_repository: Path, runner: CliRunner):
     os.chdir(sample_repository)
-    runner.invoke(init_command)
+
+    with patch.object(VenvManager, "is_initialized", return_value=True):
+        runner.invoke(init_command)
     subprocess.run(["git", "add", "."], cwd=sample_repository, check=True)
     subprocess.run(["git", "commit", "-m", "Initialize codegen"], cwd=sample_repository, check=True)
-    yield sample_repository
-    shutil.rmtree(sample_repository)
+
+    return sample_repository
diff --git a/tests/unit/codegen/cli/test_reset.py b/tests/integration/codegen/cli/test_reset.py
similarity index 100%
rename from tests/unit/codegen/cli/test_reset.py
rename to tests/integration/codegen/cli/test_reset.py
diff --git a/tests/unit/codegen/sdk/benchmark/codebase/test_codebase_reset.py b/tests/unit/codegen/sdk/benchmark/codebase/test_codebase_reset.py
index 3a8ebf689..9ac955db6 100644
--- a/tests/unit/codegen/sdk/benchmark/codebase/test_codebase_reset.py
+++ b/tests/unit/codegen/sdk/benchmark/codebase/test_codebase_reset.py
@@ -36,6 +36,7 @@ def setup():
     benchmark.pedantic(reset_codebase, setup=setup)
 
 
+@pytest.mark.skip("Skipping this test for now")
 @pytest.mark.timeout(5, func_only=True)
 @pytest.mark.parametrize("extension", ["txt", "py"])
 def test_codebase_reset_correctness(extension: str, tmp_path):
diff --git a/tests/unit/codegen/sdk/conftest.py b/tests/unit/codegen/sdk/conftest.py
index a9b12f9a5..03162cc4a 100644
--- a/tests/unit/codegen/sdk/conftest.py
+++ b/tests/unit/codegen/sdk/conftest.py
@@ -28,7 +28,7 @@ def codebase(tmp_path, original: dict[str, str], programming_language: Programmi
 
 @pytest.fixture
 def assert_expected(expected: dict[str, str], tmp_path):
-    def assert_expected(codebase: Codebase):
+    def assert_expected(codebase: Codebase) -> None:
         codebase.commit()
         for file in expected:
             assert tmp_path.joinpath(file).exists()