allow transforms to be registered in mocks

sdk/Pulumi.Tests/Mocks/TransformsTests.cs

@@ -0,0 +1,144 @@
+// Copyright 2026, Pulumi Corporation
+
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Threading;
+using System.Threading.Tasks;
+using Pulumi.Testing;
+using Xunit;
+
+namespace Pulumi.Tests.Mocks
+{
+    public class TransformsTests
+    {
+        private sealed class RecordingMocks : IMocks
+        {
+            public readonly List<ResourceTransform> StackTransforms = new List<ResourceTransform>();
+            public readonly List<InvokeTransform> InvokeTransforms = new List<InvokeTransform>();
+            public readonly List<ResourceTransform> ResourceTransforms = new List<ResourceTransform>();
+            public ImmutableDictionary<string, object>? Inputs;
+
+            public Task<(string? id, object state)> NewResourceAsync(MockResourceArgs args)
+            {
+                if (args.Name == "res")
+                {
+                    Inputs = args.Inputs;
+                    ResourceTransforms.AddRange(args.Transforms);
+                }
+                return Task.FromResult<(string?, object)>(($"{args.Name}_id", args.Inputs));
+            }
+
+            public Task<object> CallAsync(MockCallArgs args)
+                => Task.FromResult<object>(args.Args);
+
+            public Task RegisterTransform(ResourceTransform transform)
+            {
+                StackTransforms.Add(transform);
+                return Task.CompletedTask;
+            }
+
+            public Task RegisterInvokeTransform(InvokeTransform transform)
+            {
+                InvokeTransforms.Add(transform);
+                return Task.CompletedTask;
+            }
+        }
+
+        private sealed class TransformsTestResourceArgs : ResourceArgs
+        {
+            [Input("foo")]
+            public Input<string>? Foo { get; set; }
+        }
+
+        private sealed class TransformsTestResource : CustomResource
+        {
+            public TransformsTestResource(string name, TransformsTestResourceArgs args, CustomResourceOptions? options = null)
+                : base("test:index:TransformsTestResource", name, args, options)
+            {
+            }
+        }
+
+        private static Task<ResourceTransformResult?> StackTransform(ResourceTransformArgs args, CancellationToken cancellationToken)
+        {
+            var newArgs = args.Args.SetItem("foo", "stack");
+            return Task.FromResult<ResourceTransformResult?>(new ResourceTransformResult(newArgs, args.Options));
+        }
+
+        private static Task<ResourceTransformResult?> OwnTransform(ResourceTransformArgs args, CancellationToken cancellationToken)
+        {
+            var newArgs = args.Args.SetItem("foo", "own");
+            return Task.FromResult<ResourceTransformResult?>(new ResourceTransformResult(newArgs, args.Options));
+        }
+
+        private static Task<InvokeTransformResult?> AddExtraInvokeTransform(InvokeTransformArgs args, CancellationToken cancellationToken)
+        {
+            var newArgs = args.Args.SetItem("extra", "added");
+            return Task.FromResult<InvokeTransformResult?>(new InvokeTransformResult(newArgs, args.Options));
+        }
+
+        private sealed class TransformsStack : Stack
+        {
+            public TransformsStack() : base(new StackOptions
+            {
+                ResourceTransforms = { StackTransform },
+            })
+            {
+                _ = new TransformsTestResource("res", new TransformsTestResourceArgs { Foo = "orig" }, new CustomResourceOptions
+                {
+                    ResourceTransforms = { OwnTransform },
+                });
+            }
+        }
+
+        [Fact]
+        public async Task TransformsAreDeliveredButNotRun()
+        {
+            var mocks = new RecordingMocks();
+            await Deployment.TestAsync<TransformsStack>(mocks, new TestOptions { IsPreview = false });
+
+            Assert.NotNull(mocks.Inputs);
+            Assert.Equal("orig", mocks.Inputs!["foo"]);
+
+            Assert.Single(mocks.StackTransforms);
+            Assert.Single(mocks.ResourceTransforms);
+
+            var args = new ResourceTransformArgs(
+                "res",
+                "test:index:TransformsTestResource",
+                custom: true,
+                ImmutableDictionary<string, object?>.Empty.Add("foo", "orig"),
+                new CustomResourceOptions());
+
+            var ownResult = await mocks.ResourceTransforms[0](args);
+            Assert.NotNull(ownResult);
+            Assert.Equal("own", ownResult!.Value.Args["foo"]);
+
+            var stackResult = await mocks.StackTransforms[0](args);
+            Assert.NotNull(stackResult);
+            Assert.Equal("stack", stackResult!.Value.Args["foo"]);
+        }
+
+        [Fact]
+        public async Task InvokeTransformsAreDelivered()
+        {
+            var mocks = new RecordingMocks();
+            var (_, exception) = await Deployment.TryTestAsync(mocks, runner => runner.RunAsync(async () =>
+            {
+                Deployment.Instance.RegisterInvokeTransform(AddExtraInvokeTransform);
+                await ((Deployment)Deployment.InternalInstance).AwaitPendingRegistrations().ConfigureAwait(false);
+                return (IDictionary<string, object?>)new Dictionary<string, object?>();
+            }, null), new TestOptions { IsPreview = false });
+
+            Assert.Null(exception);
+            Assert.Single(mocks.InvokeTransforms);
+
+            var result = await mocks.InvokeTransforms[0](new InvokeTransformArgs(
+                "test:index:MyFunction",
+                ImmutableDictionary<string, object?>.Empty.Add("orig", "value"),
+                new InvokeOptions()));
+            Assert.NotNull(result);
+            Assert.Equal("added", result!.Value.Args["extra"]);
+            Assert.Equal("value", result.Value.Args["orig"]);
+        }
+    }
+}

sdk/Pulumi/Deployment/Deployment.cs

@@ -347,6 +347,11 @@ internal async Task RegisterInvokeTransformAsync(InvokeTransform transform)
             var callbacks = await GetCallbacksAsync(CancellationToken.None).ConfigureAwait(false);
             var callback = await AllocateInvokeTransform(callbacks.Callbacks, transform).ConfigureAwait(false);
 
+            if (Monitor is MockMonitor mockMonitor)
+            {
+                mockMonitor.RecordInvokeTransformCallback(callback.Token, transform);
+            }
+
             await Monitor.RegisterStackInvokeTransform(callback).ConfigureAwait(false);
 
             TaskCompletionSource<bool>? flushed = null;

sdk/Pulumi/Deployment/Deployment_Prepare.cs

@@ -10,6 +10,7 @@
 using Google.Protobuf;
 using Google.Protobuf.WellKnownTypes;
 using Pulumi.Serialization;
+using Pulumi.Testing;
 using Pulumirpc;
 
 namespace Pulumi
@@ -149,7 +150,12 @@ await _logger.WarnAsync(
 
                 foreach (var t in options.ResourceTransforms)
                 {
-                    transforms.Add(await AllocateTransform(callbacks.Callbacks, t).ConfigureAwait(false));
+                    var callback = await AllocateTransform(callbacks.Callbacks, t).ConfigureAwait(false);
+                    if (Monitor is MockMonitor mockMonitor)
+                    {
+                        mockMonitor.RecordTransformCallback(callback.Token, t);
+                    }
+                    transforms.Add(callback);
                 }
             }
 

sdk/Pulumi/Pulumi.xml

@@ -3655,6 +3655,22 @@
             </summary>
             <param name="args">MockRegisterResourceOutputsRequest</param>
         </member>
+        <member name="M:Pulumi.Testing.IMocks.RegisterTransform(Pulumi.ResourceTransform)">
+            <summary>
+            Invoked when the program registers a stack transform. The mock monitor does not run
+            transforms; implementations that want to exercise them can record the transform here
+            and call it from <see cref="M:Pulumi.Testing.IMocks.NewResourceAsync(Pulumi.Testing.MockResourceArgs)"/>.
+            </summary>
+            <param name="transform">The registered transform.</param>
+        </member>
+        <member name="M:Pulumi.Testing.IMocks.RegisterInvokeTransform(Pulumi.InvokeTransform)">
+            <summary>
+            Invoked when the program registers an invoke transform. The mock monitor does not run
+            transforms; implementations that want to exercise them can record the transform here
+            and call it from <see cref="M:Pulumi.Testing.IMocks.CallAsync(Pulumi.Testing.MockCallArgs)"/>.
+            </summary>
+            <param name="transform">The registered transform.</param>
+        </member>
         <member name="T:Pulumi.Testing.MockResourceArgs">
             <summary>
             MockResourceArgs for use in NewResourceAsync
@@ -3685,6 +3701,11 @@
             Resource identifier.
             </summary>
         </member>
+        <member name="P:Pulumi.Testing.MockResourceArgs.Transforms">
+            <summary>
+            The transforms declared in the resource's options. The mock monitor does not run them.
+            </summary>
+        </member>
         <member name="T:Pulumi.Testing.MockCallArgs">
             <summary>
             MockCallArgs for use in CallAsync

sdk/Pulumi/Testing/IMocks.cs

@@ -31,6 +31,22 @@ public interface IMocks
         /// </summary>
         /// <param name="args">MockRegisterResourceOutputsRequest</param>
         Task RegisterResourceOutputs(MockRegisterResourceOutputsRequest args) => Task.CompletedTask;
+
+        /// <summary>
+        /// Invoked when the program registers a stack transform. The mock monitor does not run
+        /// transforms; implementations that want to exercise them can record the transform here
+        /// and call it from <see cref="NewResourceAsync"/>.
+        /// </summary>
+        /// <param name="transform">The registered transform.</param>
+        Task RegisterTransform(ResourceTransform transform) => Task.CompletedTask;
+
+        /// <summary>
+        /// Invoked when the program registers an invoke transform. The mock monitor does not run
+        /// transforms; implementations that want to exercise them can record the transform here
+        /// and call it from <see cref="CallAsync"/>.
+        /// </summary>
+        /// <param name="transform">The registered transform.</param>
+        Task RegisterInvokeTransform(InvokeTransform transform) => Task.CompletedTask;
     }
 
     /// <summary>
@@ -62,6 +78,11 @@ public class MockResourceArgs
         /// Resource identifier.
         /// </summary>
         public string? Id { get; set; }
+
+        /// <summary>
+        /// The transforms declared in the resource's options. The mock monitor does not run them.
+        /// </summary>
+        public ImmutableArray<ResourceTransform> Transforms { get; set; } = ImmutableArray<ResourceTransform>.Empty;
     }
 
     /// <summary>

sdk/Pulumi/Testing/MockMonitor.cs

@@ -1,6 +1,7 @@
 // Copyright 2016-2020, Pulumi Corporation
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Linq;
@@ -16,6 +17,8 @@ internal class MockMonitor : IMonitor
         private readonly IMocks _mocks;
         private readonly Serializer _serializer = new Serializer(excessiveDebugOutput: false);
         private readonly Dictionary<string, object> _registeredResources = new Dictionary<string, object>();
+        private readonly ConcurrentDictionary<string, ResourceTransform> _resourceTransformCallbacks = new ConcurrentDictionary<string, ResourceTransform>();
+        private readonly ConcurrentDictionary<string, InvokeTransform> _invokeTransformCallbacks = new ConcurrentDictionary<string, InvokeTransform>();
 
         public readonly List<Resource> Resources = new List<Resource>();
 
@@ -24,6 +27,29 @@ public MockMonitor(IMocks mocks)
             _mocks = mocks;
         }
 
+        internal void RecordTransformCallback(string token, ResourceTransform transform)
+        {
+            _resourceTransformCallbacks[token] = transform;
+        }
+
+        internal void RecordInvokeTransformCallback(string token, InvokeTransform transform)
+        {
+            _invokeTransformCallbacks[token] = transform;
+        }
+
+        private ImmutableArray<ResourceTransform> ResolveTransforms(RegisterResourceRequest request)
+        {
+            var builder = ImmutableArray.CreateBuilder<ResourceTransform>();
+            foreach (var callback in request.Transforms)
+            {
+                if (_resourceTransformCallbacks.TryGetValue(callback.Token, out var transform))
+                {
+                    builder.Add(transform);
+                }
+            }
+            return builder.ToImmutable();
+        }
+
         public Task<SupportsFeatureResponse> SupportsFeatureAsync(SupportsFeatureRequest request)
         {
             // Rather than attempting to keep the list of feature flags up-to date here, we just assume the
@@ -132,6 +158,13 @@ public virtual async Task<RegisterResourceResponse> RegisterResourceAsync(Resour
 
             if (request.Type == Stack._rootPulumiStackTypeName)
             {
+                // Stack transforms are attached to the root stack's registration, so deliver
+                // them to the mocks before short-circuiting it.
+                foreach (var transform in ResolveTransforms(request))
+                {
+                    await _mocks.RegisterTransform(transform).ConfigureAwait(false);
+                }
+
                 return new RegisterResourceResponse
                 {
                     Urn = NewUrn(request.Parent, request.Type, request.Name),
@@ -146,6 +179,7 @@ public virtual async Task<RegisterResourceResponse> RegisterResourceAsync(Resour
                 Inputs = ToDictionary(request.Object),
                 Provider = request.Provider,
                 Id = request.ImportId,
+                Transforms = ResolveTransforms(request),
             }).ConfigureAwait(false);
 
             var urn = NewUrn(request.Parent, request.Type, request.Name);
@@ -225,9 +259,12 @@ private async Task<Struct> SerializeAsync(object o)
             return Serializer.CreateStruct(dict!);
         }
 
-        public Task RegisterStackInvokeTransform(Pulumirpc.Callback callback)
+        public async Task RegisterStackInvokeTransform(Pulumirpc.Callback callback)
         {
-            return Task.CompletedTask;
+            if (_invokeTransformCallbacks.TryGetValue(callback.Token, out var transform))
+            {
+                await _mocks.RegisterInvokeTransform(transform).ConfigureAwait(false);
+            }
         }
 
         public Task RegisterResourceHookAsync(RegisterResourceHookRequest request)