diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index e1859d8aa..0843fc0e4 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -9,14 +9,17 @@ jobs:
       run:
         working-directory: pgml-extension
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
-          fetch-depth: 2
+          submodules: 'recursive'
+      - name: Fetch master
+        run: |
+          git fetch origin master --depth 1
       - name: Changed files in pgml-extension
         id: pgml_extension_changed
         run: |
-          echo "PGML_EXTENSION_CHANGED_FILES=$(git diff --name-only HEAD HEAD~1 . | wc -l)" >> $GITHUB_OUTPUT
-      - name: Install dependencies
+          echo "PGML_EXTENSION_CHANGED_FILES=$(git diff --name-only HEAD origin/master . | wc -l)" >> $GITHUB_OUTPUT
+      - name: System dependencies
         if: steps.pgml_extension_changed.outputs.PGML_EXTENSION_CHANGED_FILES != '0'
         run: |
           sudo apt-get update && \
@@ -42,17 +45,8 @@ jobs:
             ~/.cargo
             pgml-extension/target
             ~/.pgrx
-          key: ${{ runner.os }}-rust-1.74-${{ hashFiles('pgml-extension/Cargo.lock') }}
-      - name: Submodules
-        if: steps.pgml_extension_changed.outputs.PGML_EXTENSION_CHANGED_FILES != '0'
-        run: |
-          git submodule update --init --recursive
-      - name: Get current version
-        id: current-version
-        run: echo "CI_BRANCH=$(git name-rev --name-only HEAD)" >> $GITHUB_OUTPUT  
-      - name: Run tests
-        env:
-          CI_BRANCH: ${{ steps.current-version.outputs.CI_BRANCH }}
+          key: ${{ runner.os }}-rust-1.74-${{ hashFiles('pgml-extension/Cargo.lock') }}-bust2
+      - name: Install pgrx
         if: steps.pgml_extension_changed.outputs.PGML_EXTENSION_CHANGED_FILES != '0'
         run: |
           curl https://sh.rustup.rs -sSf | sh -s -- -y
@@ -61,15 +55,23 @@ jobs:
 
           if [[ ! -d ~/.pgrx ]]; then
             cargo pgrx init
+            echo "shared_preload_libraries = 'pgml'" >> ~/.pgrx/data-16/postgresql.conf
           fi
-
-          git checkout master
+      - name: Update extension test
+        if: steps.pgml_extension_changed.outputs.PGML_EXTENSION_CHANGED_FILES != '0'
+        run: |
+          git checkout origin/master
           echo "\q" | cargo pgrx run
-          psql -p 28816 -h localhost -d pgml -P pager -c "CREATE EXTENSION pgml;"
-          git checkout $CI_BRANCH
+          psql -p 28816 -h localhost -d pgml -P pager -c "DROP EXTENSION IF EXISTS pgml CASCADE; DROP SCHEMA IF EXISTS pgml CASCADE; CREATE EXTENSION pgml;"
+          git checkout $GITHUB_SHA
           echo "\q" | cargo pgrx run
           psql -p 28816 -h localhost -d pgml -P pager -c "ALTER EXTENSION pgml UPDATE;"
+      - name: Unit tests
+        if: steps.pgml_extension_changed.outputs.PGML_EXTENSION_CHANGED_FILES != '0'
+        run: |
           cargo pgrx test
-#          cargo pgrx start
-#          psql -p 28815 -h 127.0.0.1 -d pgml -P pager -f tests/test.sql
-#          cargo pgrx stop
+      - name: Integration tests
+        if: steps.pgml_extension_changed.outputs.PGML_EXTENSION_CHANGED_FILES != '0'
+        run: |
+          echo "\q" | cargo pgrx run
+          psql -p 28816 -h 127.0.0.1 -d pgml -P pager -f tests/test.sql
diff --git a/pgml-extension/sql/setup_examples.sql b/pgml-extension/sql/setup_examples.sql
index 4f14924ea..2d55b54d8 100644
--- a/pgml-extension/sql/setup_examples.sql
+++ b/pgml-extension/sql/setup_examples.sql
@@ -4,7 +4,7 @@
 --- Usage:
 ---
 ---   $ cargo pgrx run --release
----   $ psql -P pager-off -h localhost -p 28813 -d pgml -f sql/setup_examples.sql
+---   $ psql -P pager-off -h localhost -p 28816 -d pgml -f sql/setup_examples.sql
 ---
 -- \set ON_ERROR_STOP true
 \timing on
diff --git a/pgml-extension/src/orm/snapshot.rs b/pgml-extension/src/orm/snapshot.rs
index 9a0c22780..402dff976 100644
--- a/pgml-extension/src/orm/snapshot.rs
+++ b/pgml-extension/src/orm/snapshot.rs
@@ -749,7 +749,9 @@ impl Snapshot {
                     self.relation_name_quoted()
                 )
             }
-            false => self.test_sampling.get_sql(&self.relation_name_quoted(), self.columns.clone()),
+            false => self
+                .test_sampling
+                .get_sql(&self.relation_name_quoted(), self.columns.clone()),
         }
     }
 
diff --git a/pgml-extension/tests/test.sql b/pgml-extension/tests/test.sql
index 1b9e3771b..a6c75dee9 100644
--- a/pgml-extension/tests/test.sql
+++ b/pgml-extension/tests/test.sql
@@ -4,7 +4,7 @@
 --- Usage:
 ---
 ---   $ cargo pgrx run --release
----   $ psql -h localhost -p 28815 -d pgml -f tests/test.sql -P pager
+---   $ psql -h localhost -p 28816 -d pgml -f tests/test.sql -P pager
 ---
 \set ON_ERROR_STOP true
 \timing on