diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 789c895..7e562be 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -33,7 +33,7 @@ jobs:
     strategy:
       matrix:
         features:
-          - dev-tunnels
+          - azure-cli-persistence
           - shell-history
     steps:
       - uses: actions/checkout@v4
diff --git a/src/azure-cli-persistence/oncreate.sh b/src/azure-cli-persistence/oncreate.sh
index 197f6cd..8526318 100644
--- a/src/azure-cli-persistence/oncreate.sh
+++ b/src/azure-cli-persistence/oncreate.sh
@@ -18,6 +18,22 @@ fi
 
 log "In OnCreate script"
 
+FEATURE_DIR="/usr/local/share/stuartleeks-devcontainer-features/azure-cli-persistence"
+LOG_FILE="$FEATURE_DIR/log.txt"
+
+log() {
+  echo "$1"
+  echo "$1" >> "$LOG_FILE"
+}
+
+if command -v sudo > /dev/null; then
+  sudo chown -R "$(id -u):$(id -g)" "$LOG_FILE"
+else
+  chown -R "$(id -u):$(id -g)" "$LOG_FILE"
+fi
+
+log "In OnCreate script"
+
 fix_permissions() {
     local dir
     dir="${1}"
@@ -47,11 +63,8 @@ old_cliextensions_folder="$HOME/.azure-old/cliextensions"
 new_cliextensions_folder="$HOME/.azure/cliextensions"
 new_cliextensions_folder_parent="$HOME/.azure"
 
-got_old_extensions_folder=false
-
 if [ "$got_old_azure_folder" = true ]; then
   if [ -d "$old_cliextensions_folder" ]; then
-    got_old_extensions_folder=true
     log "cliextensions folder found in old .azure folder"
     if [ -d "$new_cliextensions_folder" ]; then
       if [ -L "$new_cliextensions_folder" ]; then
@@ -75,13 +88,11 @@ if [ "$got_old_azure_folder" = true ]; then
       fi
     fi
   fi
-fi
-
-# If we haven't got an old .azure folder with a cliextensions folder in it, check if the new cliextensions folder is a symlink to the old one
-# And if so, remove the symlink
-# This can happen if the user has installed the azure-cli feature and specified extensions to install
-# and then later removed the extensions.
-if [ "$got_old_extensions_folder" = false ]; then
+else
+  # If we haven't got an old .azure folder with a cliextensions folder in it, check if the new cliextensions folder is a symlink to the old one
+  # And if so, remove the symlink
+  # This can happen if the user has installed the azure-cli feature and specified extensions to install
+  # and then later removed the extensions.
   if [ -L "$new_cliextensions_folder" ]; then
     symlink_target=$(readlink "$new_cliextensions_folder")
     if [ "$symlink_target" = "$old_cliextensions_folder" ]; then
diff --git a/test-project/.devcontainer/devcontainer.json b/test-project/.devcontainer/devcontainer.json
index ba3fdc6..011a171 100644
--- a/test-project/.devcontainer/devcontainer.json
+++ b/test-project/.devcontainer/devcontainer.json
@@ -14,6 +14,9 @@
 			// "extensions": "containerapp,ssh"
 		},
 		"ghcr.io/devcontainers-extra/features/fish-apt-get": {},
+		"ghcr.io/devcontainers/features/azure-cli:1": {
+			"extensions": "containerapp,ssh"
+		},
 		"./src/shell-history": {},
 		"./src/azure-cli-persistence": {}
 	}
diff --git a/test/azure-cli-persistence/_default.sh b/test/azure-cli-persistence/_default.sh
new file mode 100644
index 0000000..21c4360
--- /dev/null
+++ b/test/azure-cli-persistence/_default.sh
@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+
+set -e
+
+# Optional: Import test library bundled with the devcontainer CLI
+source dev-container-features-test-lib
+
+# Feature-specific tests
+check "cache dir permission" bash -c "test -w /dc/azure"
+check "check symlink" bash -c 'test $(readlink "$HOME/.azure") == "/dc/azure/"'
+
+# Report result
+reportResults
diff --git a/test/azure-cli-persistence/root_user.sh b/test/azure-cli-persistence/root_user.sh
new file mode 100644
index 0000000..071c03f
--- /dev/null
+++ b/test/azure-cli-persistence/root_user.sh
@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+
+set -e
+
+source _default.sh
diff --git a/test/azure-cli-persistence/scenarios.json b/test/azure-cli-persistence/scenarios.json
new file mode 100644
index 0000000..45b4e02
--- /dev/null
+++ b/test/azure-cli-persistence/scenarios.json
@@ -0,0 +1,19 @@
+{
+    "with_extensions": {
+        "image": "mcr.microsoft.com/devcontainers/base:debian",
+        "features": {
+            "ghcr.io/devcontainers/features/azure-cli:1": {
+                "extensions": "containerapp,ssh"
+            },
+            "azure-cli-persistence": {}
+        }
+    },
+    "root_user": {
+        "image": "mcr.microsoft.com/devcontainers/base:debian",
+        "features": {
+            "ghcr.io/devcontainers/features/azure-cli:1": {},
+            "azure-cli-persistence": {}
+        },
+        "remoteUser": "root"
+    }
+}
\ No newline at end of file
diff --git a/test/azure-cli-persistence/test.sh b/test/azure-cli-persistence/test.sh
index 92425a6..071c03f 100644
--- a/test/azure-cli-persistence/test.sh
+++ b/test/azure-cli-persistence/test.sh
@@ -2,11 +2,4 @@
 
 set -e
 
-# Optional: Import test library bundled with the devcontainer CLI
-source dev-container-features-test-lib
-
-# Feature-specific tests
-check "cache dir permission" bash -c "test -w /dc/azure"
-
-# Report result
-reportResults
+source _default.sh
diff --git a/test/azure-cli-persistence/with_extensions.sh b/test/azure-cli-persistence/with_extensions.sh
new file mode 100644
index 0000000..e129385
--- /dev/null
+++ b/test/azure-cli-persistence/with_extensions.sh
@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+set -e
+
+source dev-container-features-test-lib
+
+check "cache dir permission" bash -c "test -w /dc/azure"
+check "check symlink" bash -c 'test $(readlink "$HOME/.azure") == "/dc/azure/"'
+
+check "check cliextensions" bash -c 'test $(az extension list --output tsv --query "length(@)") -eq 2'
+
+reportResults
\ No newline at end of file