From 8a84a6f54a490e3842a04b15f3a181e258771892 Mon Sep 17 00:00:00 2001
From: Daniel Rojas <danielrojasc888@gmail.com>
Date: Mon, 23 Feb 2026 10:33:06 -0500
Subject: [PATCH] feat(examples): add upgradable storage contract tests

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../examples/contracts/upgradable_storage.py  | 33 +++++++++++++
 .../contracts/upgradable_storage_v2.py        | 32 +++++++++++++
 .../examples/tests/test_upgradable_storage.py | 48 +++++++++++++++++++
 3 files changed, 113 insertions(+)
 create mode 100644 tests/examples/contracts/upgradable_storage.py
 create mode 100644 tests/examples/contracts/upgradable_storage_v2.py
 create mode 100644 tests/examples/tests/test_upgradable_storage.py

diff --git a/tests/examples/contracts/upgradable_storage.py b/tests/examples/contracts/upgradable_storage.py
new file mode 100644
index 0000000..2075731
--- /dev/null
+++ b/tests/examples/contracts/upgradable_storage.py
@@ -0,0 +1,33 @@
+# v0.1.0
+# { "Depends": "py-genlayer:latest" }
+
+from genlayer import *
+
+
+class UpgradableStorage(gl.Contract):
+    storage: str
+
+    def __init__(self, initial_storage: str):
+        self.storage = initial_storage
+
+        # make the deployer an upgrader
+        root = gl.storage.Root.get()
+        root.upgraders.get().append(gl.message.sender_address)
+
+        # lock critical slots so only upgraders can modify them
+        root.lock_default()
+
+    @gl.public.view
+    def get_storage(self) -> str:
+        return self.storage
+
+    @gl.public.write
+    def update_storage(self, new_storage: str) -> None:
+        self.storage = new_storage
+
+    @gl.public.write
+    def upgrade(self, new_code: bytes) -> None:
+        root = gl.storage.Root.get()
+        code = root.code.get()
+        code.truncate()
+        code.extend(new_code)
diff --git a/tests/examples/contracts/upgradable_storage_v2.py b/tests/examples/contracts/upgradable_storage_v2.py
new file mode 100644
index 0000000..616eb57
--- /dev/null
+++ b/tests/examples/contracts/upgradable_storage_v2.py
@@ -0,0 +1,32 @@
+# v0.1.0
+# { "Depends": "py-genlayer:latest" }
+
+from genlayer import *
+
+
+class UpgradableStorage(gl.Contract):
+    # storage layout must remain compatible with v1
+    storage: str
+
+    def __init__(self):
+        pass
+
+    @gl.public.view
+    def get_storage(self) -> str:
+        return self.storage
+
+    @gl.public.write
+    def update_storage(self, new_storage: str) -> None:
+        self.storage = new_storage
+
+    # new method added in v2
+    @gl.public.view
+    def get_storage_length(self) -> int:
+        return len(self.storage)
+
+    @gl.public.write
+    def upgrade(self, new_code: bytes) -> None:
+        root = gl.storage.Root.get()
+        code = root.code.get()
+        code.truncate()
+        code.extend(new_code)
diff --git a/tests/examples/tests/test_upgradable_storage.py b/tests/examples/tests/test_upgradable_storage.py
new file mode 100644
index 0000000..d6dd40c
--- /dev/null
+++ b/tests/examples/tests/test_upgradable_storage.py
@@ -0,0 +1,48 @@
+from pathlib import Path
+from gltest import get_contract_factory
+from gltest.assertions import tx_execution_succeeded
+
+
+CONTRACTS_DIR = Path(__file__).parent.parent / "contracts"
+INITIAL_STATE = "hello"
+UPDATED_STATE = "world"
+
+
+def test_upgradable_storage():
+    # Deploy v1
+    factory = get_contract_factory(
+        contract_file_path=CONTRACTS_DIR / "upgradable_storage.py"
+    )
+    contract = factory.deploy(args=[INITIAL_STATE])
+
+    # Verify initial state
+    assert contract.get_storage(args=[]).call() == INITIAL_STATE
+
+    # Update state using v1 method
+    tx = contract.update_storage(args=[UPDATED_STATE]).transact()
+    assert tx_execution_succeeded(tx)
+    assert contract.get_storage(args=[]).call() == UPDATED_STATE
+
+    # Read v2 code and upgrade the contract
+    v2_code = (CONTRACTS_DIR / "upgradable_storage_v2.py").read_bytes()
+    tx = contract.upgrade(args=[v2_code]).transact()
+    assert tx_execution_succeeded(tx)
+
+    # After upgrade, rebuild the contract proxy from the v2 schema
+    # so we can access newly added methods
+    v2_factory = get_contract_factory(
+        contract_file_path=CONTRACTS_DIR / "upgradable_storage_v2.py"
+    )
+    contract_v2 = v2_factory.build_contract(contract_address=contract.address)
+
+    # Storage persists across upgrades
+    assert contract_v2.get_storage(args=[]).call() == UPDATED_STATE
+
+    # New v2 method works
+    assert contract_v2.get_storage_length(args=[]).call() == len(UPDATED_STATE)
+
+    # Existing methods still work
+    tx = contract_v2.update_storage(args=["upgraded"]).transact()
+    assert tx_execution_succeeded(tx)
+    assert contract_v2.get_storage(args=[]).call() == "upgraded"
+    assert contract_v2.get_storage_length(args=[]).call() == len("upgraded")