@itwin/core-backend 5.8.2 → 5.9.0-dev.10

package/lib/esm/test/schema/SchemaImportCallbacks.test.js

@@ -6,6 +6,7 @@ import { assert, expect } from "chai";
 import * as path from "path";
 import { Guid } from "@itwin/core-bentley";
 import { Code, IModel } from "@itwin/core-common";
+import { withEditTxn } from "../../EditTxn";
 import { ChannelControl, DataTransformationStrategy, IModelJsFs, StandaloneDb } from "../../core-backend";
 import { HubWrappers, IModelTestUtils } from "../IModelTestUtils";
 import { KnownTestLocations } from "../KnownTestLocations";
@@ -13,26 +14,26 @@ import { HubMock } from "../../internal/HubMock";
 describe("Schema Import Callbacks", () => {
     let imodel;
     // Test schema with version changes
-    const testSchemaV100 = () => `<?xml version="1.0" encoding="UTF-8"?>
-    <ECSchema schemaName="TestSchema" alias="ts" version="1.0.0" xmlns="http://www.bentley.com/schemas/Bentley.ECXML.3.2">
-      <ECSchemaReference name="BisCore" version="1.0.0" alias="bis"/>
-      <ECEntityClass typeName="TestElement">
-        <BaseClass>bis:DefinitionElement</BaseClass>
-        <ECProperty propertyName="StringProp" typeName="string" />
-        <ECProperty propertyName="IntProp" typeName="int" />
-        <ECProperty propertyName="ModelName" typeName="string" />
-      </ECEntityClass>
+    const testSchemaV100 = () => `<?xml version="1.0" encoding="UTF-8"?>
+    <ECSchema schemaName="TestSchema" alias="ts" version="1.0.0" xmlns="http://www.bentley.com/schemas/Bentley.ECXML.3.2">
+      <ECSchemaReference name="BisCore" version="1.0.0" alias="bis"/>
+      <ECEntityClass typeName="TestElement">
+        <BaseClass>bis:DefinitionElement</BaseClass>
+        <ECProperty propertyName="StringProp" typeName="string" />
+        <ECProperty propertyName="IntProp" typeName="int" />
+        <ECProperty propertyName="ModelName" typeName="string" />
+      </ECEntityClass>
     </ECSchema>`;
-    const testSchemaV101 = () => `<?xml version="1.0" encoding="UTF-8"?>
-    <ECSchema schemaName="TestSchema" alias="ts" version="1.0.1" xmlns="http://www.bentley.com/schemas/Bentley.ECXML.3.2">
-      <ECSchemaReference name="BisCore" version="1.0.0" alias="bis"/>
-      <ECEntityClass typeName="TestElement">
-        <BaseClass>bis:DefinitionElement</BaseClass>
-        <ECProperty propertyName="StringProp" typeName="string" />
-        <ECProperty propertyName="IntProp" typeName="int" />
-        <ECProperty propertyName="NewProp" typeName="string" />
-        <ECProperty propertyName="ModelName" typeName="string" />
-      </ECEntityClass>
+    const testSchemaV101 = () => `<?xml version="1.0" encoding="UTF-8"?>
+    <ECSchema schemaName="TestSchema" alias="ts" version="1.0.1" xmlns="http://www.bentley.com/schemas/Bentley.ECXML.3.2">
+      <ECSchemaReference name="BisCore" version="1.0.0" alias="bis"/>
+      <ECEntityClass typeName="TestElement">
+        <BaseClass>bis:DefinitionElement</BaseClass>
+        <ECProperty propertyName="StringProp" typeName="string" />
+        <ECProperty propertyName="IntProp" typeName="int" />
+        <ECProperty propertyName="NewProp" typeName="string" />
+        <ECProperty propertyName="ModelName" typeName="string" />
+      </ECEntityClass>
     </ECSchema>`;
     beforeEach(() => {
         const testFileName = IModelTestUtils.prepareOutputFile("SchemaImportCallbacks", `SchemaCallbackTest_${Guid.createValue()}.bim`);
@@ -122,37 +123,39 @@ describe("Schema Import Callbacks", () => {
                 stringProp: "original value of first element",
                 intProp: 42,
             };
-            const elementId1 = imodel.elements.insertElement(elementProps);
+            const elementId1 = withEditTxn(imodel, (txn) => txn.insertElement(elementProps));
             const elementIds = [elementId1];
-            await imodel.importSchemaStrings([testSchemaV101()], {
-                schemaImportCallbacks: {
-                    preSchemaImportCallback: async (context) => {
-                        // Create another element before the schema import
-                        assert.equal(elementIds.length, 1);
-                        const element = context.iModel.elements.getElementProps(elementIds[0]);
-                        assert.isDefined(element.stringProp);
-                        assert.isDefined(element.intProp);
-                        elementProps.stringProp = "original value of second element";
-                        elementProps.intProp = 84;
-                        elementIds.push(imodel.elements.insertElement(elementProps));
-                        const cached = { ids: elementIds };
-                        return {
-                            transformStrategy: DataTransformationStrategy.InMemory,
-                            cachedData: cached,
-                        };
-                    },
-                    postSchemaImportCallback: async (context) => {
-                        // Use cached data to update element with new property
-                        assert.isDefined(context.resources.cachedData?.ids);
-                        const updatedElementProps = context.iModel.elements.getElementProps(context.resources.cachedData.ids[1]);
-                        assert.isDefined(updatedElementProps.stringProp);
-                        assert.isDefined(updatedElementProps.intProp);
-                        assert.isUndefined(updatedElementProps.newProp);
-                        updatedElementProps.stringProp = "modified in postImport";
-                        updatedElementProps.newProp = `New Prop Added`;
-                        context.iModel.elements.updateElement(updatedElementProps);
+            await withEditTxn(imodel, "cache callback writes", async (txn) => {
+                await imodel.importSchemaStrings([testSchemaV101()], {
+                    schemaImportCallbacks: {
+                        preSchemaImportCallback: async (context) => {
+                            // Create another element before the schema import
+                            assert.equal(elementIds.length, 1);
+                            const element = context.iModel.elements.getElementProps(elementIds[0]);
+                            assert.isDefined(element.stringProp);
+                            assert.isDefined(element.intProp);
+                            elementProps.stringProp = "original value of second element";
+                            elementProps.intProp = 84;
+                            elementIds.push(txn.insertElement(elementProps));
+                            const cached = { ids: elementIds };
+                            return {
+                                transformStrategy: DataTransformationStrategy.InMemory,
+                                cachedData: cached,
+                            };
+                        },
+                        postSchemaImportCallback: async (context) => {
+                            // Use cached data to update element with new property
+                            assert.isDefined(context.resources.cachedData?.ids);
+                            const updatedElementProps = context.iModel.elements.getElementProps(context.resources.cachedData.ids[1]);
+                            assert.isDefined(updatedElementProps.stringProp);
+                            assert.isDefined(updatedElementProps.intProp);
+                            assert.isUndefined(updatedElementProps.newProp);
+                            updatedElementProps.stringProp = "modified in postImport";
+                            updatedElementProps.newProp = `New Prop Added`;
+                            txn.updateElement(updatedElementProps);
+                        },
                     },
-                },
+                });
             });
             // Verify the transformation
             let finalElementProps = imodel.elements.getElementProps(elementIds[0]);
@@ -192,26 +195,27 @@ describe("Schema Import Callbacks", () => {
                 stringProp: "snapshot test",
                 intProp: 123,
             };
-            const elementId = imodel.elements.insertElement(elementProps);
-            imodel.saveChanges("Insert element before schema upgrade");
+            const elementId = withEditTxn(imodel, "Insert element before schema upgrade", (txn) => txn.insertElement(elementProps));
             // Import updated schema with snapshot strategy
             let originalStringValue;
-            await imodel.importSchemaStrings([testSchemaV101()], {
-                schemaImportCallbacks: {
-                    preSchemaImportCallback: async () => ({ transformStrategy: DataTransformationStrategy.Snapshot }),
-                    postSchemaImportCallback: async (context) => {
-                        assert.isDefined(context.resources.snapshot);
-                        assert.equal(context.resources.snapshot?.getSchemaProps("TestSchema").version, "01.00.00");
-                        assert.equal(imodel.getSchemaProps("TestSchema").version, "01.00.01");
-                        // Read original value from snapshot
-                        const snapshotElementProps = context.resources.snapshot.elements.getElementProps(elementId);
-                        originalStringValue = snapshotElementProps.stringProp;
-                        // Update element in main iModel with new property based on snapshot data
-                        const updatedElementProps = context.iModel.elements.getElementProps(elementId);
-                        updatedElementProps.newProp = `Original was: ${originalStringValue}`;
-                        context.iModel.elements.updateElement(updatedElementProps);
+            await withEditTxn(imodel, "snapshot callback writes", async (txn) => {
+                await imodel.importSchemaStrings([testSchemaV101()], {
+                    schemaImportCallbacks: {
+                        preSchemaImportCallback: async () => ({ transformStrategy: DataTransformationStrategy.Snapshot }),
+                        postSchemaImportCallback: async (context) => {
+                            assert.isDefined(context.resources.snapshot);
+                            assert.equal(context.resources.snapshot?.getSchemaProps("TestSchema").version, "01.00.00");
+                            assert.equal(imodel.getSchemaProps("TestSchema").version, "01.00.01");
+                            // Read original value from snapshot
+                            const snapshotElementProps = context.resources.snapshot.elements.getElementProps(elementId);
+                            originalStringValue = snapshotElementProps.stringProp;
+                            // Update element in main iModel with new property based on snapshot data
+                            const updatedElementProps = context.iModel.elements.getElementProps(elementId);
+                            updatedElementProps.newProp = `Original was: ${originalStringValue}`;
+                            txn.updateElement(updatedElementProps);
+                        },
                     },
-                },
+                });
             });
             assert.equal(originalStringValue, "snapshot test");
             const finalElement = imodel.elements.getElement(elementId);
@@ -264,23 +268,24 @@ describe("Schema Import Callbacks", () => {
                 intProp: 1,
             };
             assert.equal(imodel.getSchemaProps("TestSchema").version, "01.00.00");
-            const elementId = imodel.elements.insertElement(elementProps);
-            imodel.saveChanges("Insert test element");
+            const elementId = withEditTxn(imodel, "Insert test element", (txn) => txn.insertElement(elementProps));
             // Try to import with failing callback
             try {
-                await imodel.importSchemaStrings([testSchemaV101()], {
-                    schemaImportCallbacks: {
-                        postSchemaImportCallback: async (context) => {
-                            // Make a change
-                            const updatedElementProps = context.iModel.elements.getElementProps(elementId);
-                            updatedElementProps.intProp += 1;
-                            updatedElementProps.stringProp = "should be reverted";
-                            updatedElementProps.newProp = "should be reverted";
-                            context.iModel.elements.updateElement(updatedElementProps);
-                            // Then throw error
-                            throw new Error("Intentional callback failure");
+                await withEditTxn(imodel, "failing callback writes", async (txn) => {
+                    await imodel.importSchemaStrings([testSchemaV101()], {
+                        schemaImportCallbacks: {
+                            postSchemaImportCallback: async (context) => {
+                                // Make a change
+                                const updatedElementProps = context.iModel.elements.getElementProps(elementId);
+                                updatedElementProps.intProp += 1;
+                                updatedElementProps.stringProp = "should be reverted";
+                                updatedElementProps.newProp = "should be reverted";
+                                txn.updateElement(updatedElementProps);
+                                // Then throw error
+                                throw new Error("Intentional callback failure");
+                            },
                         },
-                    },
+                    });
                 });
                 assert.fail("Should have thrown error");
             }
@@ -440,22 +445,23 @@ describe("Schema Import Callbacks", () => {
                 stringProp: "test element",
                 intProp: 100,
             };
-            const elementId = imodel.elements.insertElement(elementProps);
-            imodel.saveChanges("Create test element");
+            const elementId = withEditTxn(imodel, "Create test element", (txn) => txn.insertElement(elementProps));
             // Now REMOVE the shared channel permission
             imodel.channels.removeAllowedChannel("shared");
             // Try to import schema and modify element - should fail
             try {
-                await imodel.importSchemaStrings([testSchemaV101()], {
-                    schemaImportCallbacks: {
-                        preSchemaImportCallback: async (context) => {
-                            // This should throw because shared channel is not allowed
-                            const updatedProps = context.iModel.elements.getElementProps(elementId);
-                            updatedProps.newProp = "This should fail";
-                            context.iModel.elements.updateElement(updatedProps); // Should throw here
-                            return { transformStrategy: DataTransformationStrategy.None };
-                        }
-                    },
+                await withEditTxn(imodel, "channel validation pre import", async (txn) => {
+                    await imodel.importSchemaStrings([testSchemaV101()], {
+                        schemaImportCallbacks: {
+                            preSchemaImportCallback: async (context) => {
+                                // This should throw because shared channel is not allowed
+                                const updatedProps = context.iModel.elements.getElementProps(elementId);
+                                updatedProps.newProp = "This should fail";
+                                txn.updateElement(updatedProps); // Should throw here
+                                return { transformStrategy: DataTransformationStrategy.None };
+                            }
+                        },
+                    });
                 });
                 assert.fail("Should have thrown ChannelConstraintViolation");
             }
@@ -474,22 +480,23 @@ describe("Schema Import Callbacks", () => {
                 stringProp: "test element",
                 intProp: 100,
             };
-            const elementId = imodel.elements.insertElement(elementProps);
-            imodel.saveChanges("Create test element");
+            const elementId = withEditTxn(imodel, "Create test element", (txn) => txn.insertElement(elementProps));
             // Now REMOVE the shared channel permission
             imodel.channels.removeAllowedChannel("shared");
             // Try to import schema and modify element - should fail
             try {
-                await imodel.importSchemaStrings([testSchemaV101()], {
-                    schemaImportCallbacks: {
-                        preSchemaImportCallback: async () => ({ transformStrategy: DataTransformationStrategy.None }),
-                        postSchemaImportCallback: async (context) => {
-                            // This should throw because shared channel is not allowed
-                            const updatedProps = context.iModel.elements.getElementProps(elementId);
-                            updatedProps.newProp = "This should fail";
-                            context.iModel.elements.updateElement(updatedProps); // Should throw here
+                await withEditTxn(imodel, "channel validation post import", async (txn) => {
+                    await imodel.importSchemaStrings([testSchemaV101()], {
+                        schemaImportCallbacks: {
+                            preSchemaImportCallback: async () => ({ transformStrategy: DataTransformationStrategy.None }),
+                            postSchemaImportCallback: async (context) => {
+                                // This should throw because shared channel is not allowed
+                                const updatedProps = context.iModel.elements.getElementProps(elementId);
+                                updatedProps.newProp = "This should fail";
+                                txn.updateElement(updatedProps); // Should throw here
+                            },
                         },
-                    },
+                    });
                 });
                 assert.fail("Should have thrown ChannelConstraintViolation");
             }
@@ -514,24 +521,25 @@ describe("Schema Import Callbacks", () => {
                 stringProp: "snapshot test",
                 intProp: 123,
             };
-            const elementId = imodel.elements.insertElement(elementProps);
-            imodel.saveChanges("Create test element");
+            const elementId = withEditTxn(imodel, "Create test element", (txn) => txn.insertElement(elementProps));
             // Disable channel before schema import with snapshot
             imodel.channels.removeAllowedChannel("shared");
             try {
-                await imodel.importSchemaStrings([testSchemaV101()], {
-                    schemaImportCallbacks: {
-                        preSchemaImportCallback: async () => ({ transformStrategy: DataTransformationStrategy.Snapshot }),
-                        postSchemaImportCallback: async (context) => {
-                            // Can read from snapshot (it's read-only, no channel check)
-                            const snapshotProps = context.resources.snapshot.elements.getElementProps(elementId);
-                            assert.equal(snapshotProps.stringProp, "snapshot test");
-                            // But can't modify in main iModel without channel permission
-                            const updatedProps = context.iModel.elements.getElementProps(elementId);
-                            updatedProps.newProp = "This should fail";
-                            context.iModel.elements.updateElement(updatedProps); // Should throw
+                await withEditTxn(imodel, "channel validation snapshot import", async (txn) => {
+                    await imodel.importSchemaStrings([testSchemaV101()], {
+                        schemaImportCallbacks: {
+                            preSchemaImportCallback: async () => ({ transformStrategy: DataTransformationStrategy.Snapshot }),
+                            postSchemaImportCallback: async (context) => {
+                                // Can read from snapshot (it's read-only, no channel check)
+                                const snapshotProps = context.resources.snapshot.elements.getElementProps(elementId);
+                                assert.equal(snapshotProps.stringProp, "snapshot test");
+                                // But can't modify in main iModel without channel permission
+                                const updatedProps = context.iModel.elements.getElementProps(elementId);
+                                updatedProps.newProp = "This should fail";
+                                txn.updateElement(updatedProps); // Should throw
+                            },
                         },
-                    },
+                    });
                 });
                 assert.fail("Should have thrown ChannelConstraintViolation");
             }
@@ -678,16 +686,17 @@ describe("Schema Import Callbacks", () => {
                 return undefined;
             return aspects[0].asAny.version;
         }
-        function setChannelVersion(iModel, version) {
+        function setChannelVersion(txn, version) {
+            const iModel = txn.iModel;
             const aspects = iModel.elements.getAspects(channelRootId, "BisCore:ChannelRootAspect");
             assert.equal(aspects.length, 1, "Should have exactly one ChannelRootAspect");
             const aspect = aspects[0];
             aspect.asAny.version = version;
-            iModel.elements.updateAspect(aspect.toJSON());
+            txn.updateAspect(aspect.toJSON());
         }
         // Code that simulates a channel upgrade.
         // This simulates elements being moved to different models as per the new channel organization.
-        const channelUpgradeCallback = async (context) => {
+        const channelUpgradeCallback = (txn) => async (context) => {
             context.data.elementIds.forEach((id) => {
                 const elementProps = context.iModel.elements.getElementProps(id);
                 if (elementProps.stringProp === "Material1") {
@@ -699,11 +708,10 @@ describe("Schema Import Callbacks", () => {
                 else if (elementProps.stringProp === "Category1") {
                     elementProps.modelName = "StyleModel";
                 }
-                context.iModel.elements.updateElement(elementProps);
+                txn.updateElement(elementProps);
             });
-            setChannelVersion(context.iModel, "1.0.1");
+            setChannelVersion(txn, "1.0.1");
             assert.equal(getChannelVersion(context.iModel), "1.0.1");
-            context.iModel.saveChanges();
         };
         it("need for channel reorganization before schema import", async () => {
             // Initial setup: Import v1.0.0 schema
@@ -712,13 +720,14 @@ describe("Schema Import Callbacks", () => {
             const channelKey = "TestChannel";
             imodel.channels.addAllowedChannel(channelKey);
             // Create a channel
-            channelRootId = imodel.channels.insertChannelSubject({
+            channelRootId = withEditTxn(imodel, (txn) => imodel.channels.insertChannelSubject({
                 subjectName: "Test Channel",
                 channelKey,
-            });
+                txn,
+            }));
             assert.isUndefined(getChannelVersion(imodel), "Version should be undefined initially");
             // Set version to 1.0.0
-            setChannelVersion(imodel, "1.0.0");
+            withEditTxn(imodel, (txn) => setChannelVersion(txn, "1.0.0"));
             // Enable shared channel and create elements
             imodel.channels.addAllowedChannel("shared");
             const rootModel = imodel.models.getModel(IModel.dictionaryId);
@@ -734,40 +743,41 @@ describe("Schema Import Callbacks", () => {
                     intProp: 100 + index * 100,
                     modelName: "DefinitionModel"
                 };
-                elementIds.push(imodel.elements.insertElement(elementProps));
+                elementIds.push(withEditTxn(imodel, (txn) => txn.insertElement(elementProps)));
             }
-            imodel.saveChanges("Create elements in a single model under the root channel structure");
             // Scenario: Let's assume schema has evolved and v1.0.1 expects elements to be in their dedicated models (i.e. it relies on channel v1.0.1)
             // The post schema upgrade code will look for elements in their own models.
             // If the channel is still in version 1.0.0, the callback will obviously fail.
-            await imodel.importSchemaStrings([testSchemaV101()], {
-                data: { elementIds },
-                channelUpgrade: {
-                    channelKey,
-                    fromVersion: "1.0.0",
-                    toVersion: "1.0.1",
-                    callback: channelUpgradeCallback,
-                },
-                schemaImportCallbacks: {
-                    preSchemaImportCallback: async () => {
-                        return {
-                            transformStrategy: DataTransformationStrategy.InMemory,
-                            cachedData: {
-                                elementIds,
-                            },
-                        };
+            await withEditTxn(imodel, "channel reorganization callback writes", async (txn) => {
+                await imodel.importSchemaStrings([testSchemaV101()], {
+                    data: { elementIds },
+                    channelUpgrade: {
+                        channelKey,
+                        fromVersion: "1.0.0",
+                        toVersion: "1.0.1",
+                        callback: channelUpgradeCallback(txn),
                     },
-                    postSchemaImportCallback: async (context) => {
-                        // Now schema v1.0.1 expects elements to be in their dedicated models.
-                        const elementIdList = context.resources.cachedData.elementIds;
-                        assert.equal(imodel.getSchemaProps("TestSchema").version, "01.00.01");
-                        assert.equal(getChannelVersion(context.iModel), "1.0.1");
-                        elementIdList.forEach((elementId) => {
-                            const elementProps = context.iModel.elements.getElementProps(elementId);
-                            assert.notEqual(elementProps.modelName, "DefinitionModel", "Element should have been moved to new model as part of the channel upgrade");
-                        });
+                    schemaImportCallbacks: {
+                        preSchemaImportCallback: async () => {
+                            return {
+                                transformStrategy: DataTransformationStrategy.InMemory,
+                                cachedData: {
+                                    elementIds,
+                                },
+                            };
+                        },
+                        postSchemaImportCallback: async (context) => {
+                            // Now schema v1.0.1 expects elements to be in their dedicated models.
+                            const elementIdList = context.resources.cachedData.elementIds;
+                            assert.equal(imodel.getSchemaProps("TestSchema").version, "01.00.01");
+                            assert.equal(getChannelVersion(context.iModel), "1.0.1");
+                            elementIdList.forEach((elementId) => {
+                                const elementProps = context.iModel.elements.getElementProps(elementId);
+                                assert.notEqual(elementProps.modelName, "DefinitionModel", "Element should have been moved to new model as part of the channel upgrade");
+                            });
+                        },
                     },
-                },
+                });
             });
         });
     });
@@ -787,10 +797,11 @@ describe("Schema Import Callbacks", () => {
             await briefcaseDb.importSchemaStrings([testSchemaV100()]);
             briefcaseDb.channels.addAllowedChannel(channelKey);
             if (briefcaseDb.channels.queryChannelRoot(channelKey) === undefined) {
-                briefcaseDb.channels.insertChannelSubject({
+                withEditTxn(briefcaseDb, (txn) => briefcaseDb.channels.insertChannelSubject({
                     subjectName: "Test Channel",
                     channelKey,
-                });
+                    txn,
+                }));
             }
             const model = briefcaseDb.models.getModel(IModel.dictionaryId);
             const elementProps = {
@@ -800,8 +811,7 @@ describe("Schema Import Callbacks", () => {
                 stringProp: "test",
                 intProp: 100,
             };
-            const elementId = briefcaseDb.elements.insertElement(elementProps);
-            briefcaseDb.saveChanges("Create test element");
+            const elementId = withEditTxn(briefcaseDb, "Create test element", (txn) => txn.insertElement(elementProps));
             await briefcaseDb.pushChanges({ description: "Create test element" });
             return [briefcaseDb, elementId];
         }
@@ -809,41 +819,42 @@ describe("Schema Import Callbacks", () => {
             const [briefcaseDb, elementId] = await setupBriefcase();
             // Try to modify without acquiring locks in the callback
             try {
-                await briefcaseDb.importSchemaStrings([testSchemaV101()], {
-                    channelUpgrade: {
-                        channelKey,
-                        fromVersion: "1.0.0",
-                        toVersion: "1.0.1",
-                        callback: async (context) => {
-                            // Intentionally NOT acquiring locks
-                            const props = context.iModel.elements.getElementProps(elementId);
-                            props.stringProp = "should fail";
-                            context.iModel.elements.updateElement(props);
-                            context.iModel.saveChanges("Should fail");
+                await withEditTxn(briefcaseDb, "channel upgrade without locks", async (txn) => {
+                    await briefcaseDb.importSchemaStrings([testSchemaV101()], {
+                        channelUpgrade: {
+                            channelKey,
+                            fromVersion: "1.0.0",
+                            toVersion: "1.0.1",
+                            callback: async (context) => {
+                                // Intentionally NOT acquiring locks
+                                const props = context.iModel.elements.getElementProps(elementId);
+                                props.stringProp = "should fail";
+                                txn.updateElement(props);
+                            },
                         },
-                    },
+                    });
                 });
                 assert.fail("Should have thrown error about missing locks");
             }
             catch (err) {
                 assert.equal(err.iTwinErrorId.key.name, "Lock Not Held");
-                briefcaseDb.abandonChanges();
             }
             // Try again with locks
             try {
-                await briefcaseDb.importSchemaStrings([testSchemaV101()], {
-                    channelUpgrade: {
-                        channelKey,
-                        fromVersion: "1.0.0",
-                        toVersion: "1.0.1",
-                        callback: async (context) => {
-                            await context.iModel.locks.acquireLocks({ exclusive: elementId });
-                            const props = context.iModel.elements.getElementProps(elementId);
-                            props.stringProp = "updated with locks";
-                            context.iModel.elements.updateElement(props);
-                            context.iModel.saveChanges("Updated with locks");
+                await withEditTxn(briefcaseDb, "channel upgrade with locks", async (txn) => {
+                    await briefcaseDb.importSchemaStrings([testSchemaV101()], {
+                        channelUpgrade: {
+                            channelKey,
+                            fromVersion: "1.0.0",
+                            toVersion: "1.0.1",
+                            callback: async (context) => {
+                                await context.iModel.locks.acquireLocks({ exclusive: elementId });
+                                const props = context.iModel.elements.getElementProps(elementId);
+                                props.stringProp = "updated with locks";
+                                txn.updateElement(props);
+                            },
                         },
-                    },
+                    });
                 });
             }
             catch (err) {
@@ -855,36 +866,38 @@ describe("Schema Import Callbacks", () => {
             const [briefcaseDb, elementId] = await setupBriefcase();
             // Try to modify without acquiring locks in the callback
             try {
-                await briefcaseDb.importSchemaStrings([testSchemaV101()], {
-                    schemaImportCallbacks: {
-                        preSchemaImportCallback: async (context) => {
-                            // Intentionally NOT acquiring locks
-                            const props = context.iModel.elements.getElementProps(elementId);
-                            props.stringProp = "should fail";
-                            context.iModel.elements.updateElement(props);
-                            return { transformStrategy: DataTransformationStrategy.None };
+                await withEditTxn(briefcaseDb, "pre import without locks", async (txn) => {
+                    await briefcaseDb.importSchemaStrings([testSchemaV101()], {
+                        schemaImportCallbacks: {
+                            preSchemaImportCallback: async (context) => {
+                                // Intentionally NOT acquiring locks
+                                const props = context.iModel.elements.getElementProps(elementId);
+                                props.stringProp = "should fail";
+                                txn.updateElement(props);
+                                return { transformStrategy: DataTransformationStrategy.None };
+                            },
                         },
-                    },
+                    });
                 });
                 assert.fail("Should have thrown error about missing locks");
             }
             catch (err) {
                 assert.equal(err.name, "Lock Not Held");
-                briefcaseDb.abandonChanges();
             }
             // Try again with locks
             try {
-                await briefcaseDb.importSchemaStrings([testSchemaV101()], {
-                    schemaImportCallbacks: {
-                        preSchemaImportCallback: async (context) => {
-                            await context.iModel.locks.acquireLocks({ exclusive: elementId });
-                            const props = context.iModel.elements.getElementProps(elementId);
-                            props.stringProp = "should fail";
-                            context.iModel.elements.updateElement(props);
-                            context.iModel.saveChanges("Should fail");
-                            return { transformStrategy: DataTransformationStrategy.None };
+                await withEditTxn(briefcaseDb, "pre import with locks", async (txn) => {
+                    await briefcaseDb.importSchemaStrings([testSchemaV101()], {
+                        schemaImportCallbacks: {
+                            preSchemaImportCallback: async (context) => {
+                                await context.iModel.locks.acquireLocks({ exclusive: elementId });
+                                const props = context.iModel.elements.getElementProps(elementId);
+                                props.stringProp = "should fail";
+                                txn.updateElement(props);
+                                return { transformStrategy: DataTransformationStrategy.None };
+                            },
                         },
-                    },
+                    });
                 });
             }
             catch (err) {
@@ -895,16 +908,18 @@ describe("Schema Import Callbacks", () => {
         it("lock acquisition is not required during postSchemaImportCallback", async () => {
             const [briefcaseDb, elementId] = await setupBriefcase();
             let postImportCalled = false;
-            await briefcaseDb.importSchemaStrings([testSchemaV101()], {
-                schemaImportCallbacks: {
-                    postSchemaImportCallback: async (context) => {
-                        postImportCalled = true;
-                        // Schema lock is already held at this point
-                        const props = context.iModel.elements.getElementProps(elementId);
-                        props.newProp = "added in postImport with schema lock held";
-                        context.iModel.elements.updateElement(props);
+            await withEditTxn(briefcaseDb, "post import with schema lock", async (txn) => {
+                await briefcaseDb.importSchemaStrings([testSchemaV101()], {
+                    schemaImportCallbacks: {
+                        postSchemaImportCallback: async (context) => {
+                            postImportCalled = true;
+                            // Schema lock is already held at this point
+                            const props = context.iModel.elements.getElementProps(elementId);
+                            props.newProp = "added in postImport with schema lock held";
+                            txn.updateElement(props);
+                        },
                     },
-                },
+                });
             });
             assert.isTrue(postImportCalled, "Post-import callback should have been called");
             const finalProps = briefcaseDb.elements.getElementProps(elementId);