Fix WithReference to only use IResourceWithEndpoints

src/Aspire.Hosting/ResourceBuilderExtensions.cs

@@ -481,7 +481,7 @@ private static Action<EnvironmentCallbackContext> CreateEndpointReferenceEnviron
                     context.EnvironmentVariables[$"{EnvironmentVariableNameEncoder.Encode(serviceKey)}_{encodedEndpointName.ToUpperInvariant()}"] = endpoint;
                 }
 
-                if (flags.HasFlag(ReferenceEnvironmentInjectionFlags.ServiceDiscovery))
+                if (flags.HasFlag(ReferenceEnvironmentInjectionFlags.ServiceDiscovery) && annotation.Resource is IResourceWithServiceDiscovery)
                 {
                     context.EnvironmentVariables[$"services__{serviceName}__{endpointName}__0"] = endpoint;
                 }
@@ -639,6 +639,46 @@ public static IResourceBuilder<TDestination> WithReference<TDestination>(this IR
         return builder;
     }
 
+    /// <summary>
+    /// Injects endpoint information as environment variables from the project resource into the destination resource, using the source resource's name as the service name.
+    /// Each endpoint defined on the project resource will be injected using the format defined by the <see cref="ReferenceEnvironmentInjectionAnnotation"/> on the destination resource, i.e.
+    /// either "services__{sourceResourceName}__{endpointName}__{endpointIndex}={uriString}" for .NET service discovery, or "{RESOURCE_ENDPOINT}={uri}" for endpoint injection.
+    /// </summary>
+    /// <typeparam name="TDestination">The destination resource.</typeparam>
+    /// <param name="builder">The resource where the endpoint information will be injected.</param>
+    /// <param name="source">The resource from which to extract endpoint information.</param>
+    /// <returns>The <see cref="IResourceBuilder{T}"/>.</returns>
+    [AspireExport("withEndpoints", Description = "Adds all endpoint references to another resource")]
+    public static IResourceBuilder<TDestination> WithEndpoints<TDestination>(this IResourceBuilder<TDestination> builder, IResourceBuilder<IResourceWithEndpoints> source)
+        where TDestination : IResourceWithEnvironment
+    {
+        ArgumentNullException.ThrowIfNull(builder);
+        ArgumentNullException.ThrowIfNull(source);
+
+        ApplyEndpoints(builder, source.Resource);
+        return builder;
+    }
+
+    /// <summary>
+    /// Injects endpoint information as environment variables from the project resource into the destination resource, using the source resource's name as the service name.
+    /// Each endpoint defined on the project resource will be injected using the format defined by the <see cref="ReferenceEnvironmentInjectionAnnotation"/> on the destination resource, i.e.
+    /// either "services__{name}__{endpointName}__{endpointIndex}={uriString}" for .NET service discovery, or "{name}_{ENDPOINT}={uri}" for endpoint injection.
+    /// </summary>
+    /// <typeparam name="TDestination">The destination resource.</typeparam>
+    /// <param name="builder">The resource where the endpoint information will be injected.</param>
+    /// <param name="source">The resource from which to extract endpoint information.</param>
+    /// <param name="name">The name of the resource for the environment variable.</param>
+    /// <returns>The <see cref="IResourceBuilder{T}"/>.</returns>
+    public static IResourceBuilder<TDestination> WithEndpoints<TDestination>(this IResourceBuilder<TDestination> builder, IResourceBuilder<IResourceWithEndpoints> source, string name)
+        where TDestination : IResourceWithEnvironment
+    {
+        ArgumentNullException.ThrowIfNull(builder);
+        ArgumentNullException.ThrowIfNull(source);
+
+        ApplyEndpoints(builder, source.Resource, endpointName: null, name);
+        return builder;
+    }
+
     /// <summary>
     /// Injects service discovery and endpoint information as environment variables from the uri into the destination resource, using the name as the service name.
     /// The uri will be injected using the format defined by the <see cref="ReferenceEnvironmentInjectionAnnotation"/> on the destination resource, i.e.

tests/Aspire.Hosting.Tests/WithEndpointsTests.cs

@@ -0,0 +1,170 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Aspire.Hosting.Tests.Utils;
+using Aspire.Hosting.Utils;
+using Microsoft.AspNetCore.InternalTesting;
+
+namespace Aspire.Hosting.Tests;
+
+public class WithEndpointsTests
+{
+    [Fact]
+    public async Task ResourceNamesWithDashesAreEncodedInEnvironmentVariables()
+    {
+        using var builder = TestDistributedApplicationBuilder.Create();
+
+        var projectA = builder.AddProject<ProjectA>("project-a")
+                .WithHttpsEndpoint(1000, 2000, "mybinding")
+                .WithEndpoint("mybinding", e => e.AllocatedEndpoint = new AllocatedEndpoint(e, "localhost", 2000));
+
+        var projectB = builder.AddProject<ProjectB>("consumer")
+            .WithEndpoints(projectA);
+
+        var config = await EnvironmentVariableEvaluator.GetEnvironmentVariablesAsync(projectB.Resource, DistributedApplicationOperation.Run, TestServiceProvider.Instance).DefaultTimeout();
+
+        Assert.Equal("https://localhost:2000", config["services__project-a__mybinding__0"]);
+        Assert.Equal("https://localhost:2000", config["PROJECT_A_MYBINDING"]);
+        Assert.DoesNotContain("services__project_a__mybinding__0", config.Keys);
+        Assert.DoesNotContain("PROJECT-A_MYBINDING", config.Keys);
+    }
+
+    [Fact]
+    public async Task OverriddenServiceNamesAreEncodedInEnvironmentVariables()
+    {
+        using var builder = TestDistributedApplicationBuilder.Create();
+
+        var projectA = builder.AddProject<ProjectA>("project-a")
+                .WithHttpsEndpoint(1000, 2000, "mybinding")
+                .WithEndpoint("mybinding", e => e.AllocatedEndpoint = new AllocatedEndpoint(e, "localhost", 2000));
+
+        var projectB = builder.AddProject<ProjectB>("consumer")
+            .WithEndpoints(projectA, "custom-name");
+
+        var config = await EnvironmentVariableEvaluator.GetEnvironmentVariablesAsync(projectB.Resource, DistributedApplicationOperation.Run, TestServiceProvider.Instance).DefaultTimeout();
+
+        Assert.Equal("https://localhost:2000", config["services__custom-name__mybinding__0"]);
+        Assert.Equal("https://localhost:2000", config["custom_name_MYBINDING"]);
+        Assert.DoesNotContain("services__custom_name__mybinding__0", config.Keys);
+        Assert.DoesNotContain("custom-name_MYBINDING", config.Keys);
+    }
+
+    [Theory]
+    [InlineData(ReferenceEnvironmentInjectionFlags.All)]
+    [InlineData(ReferenceEnvironmentInjectionFlags.ConnectionProperties)]
+    [InlineData(ReferenceEnvironmentInjectionFlags.ConnectionString)]
+    [InlineData(ReferenceEnvironmentInjectionFlags.ServiceDiscovery)]
+    [InlineData(ReferenceEnvironmentInjectionFlags.Endpoints)]
+    [InlineData(ReferenceEnvironmentInjectionFlags.None)]
+    public async Task ProjectWithEndpointRespectsCustomEnvironmentVariableNaming(ReferenceEnvironmentInjectionFlags flags)
+    {
+        using var builder = TestDistributedApplicationBuilder.Create();
+
+        // Create a binding and its matching annotation (simulating DCP behavior)
+        var projectA = builder.AddProject<ProjectA>("projecta")
+                .WithHttpsEndpoint(1000, 2000, "mybinding")
+                .WithEndpoint("mybinding", e => e.AllocatedEndpoint = new AllocatedEndpoint(e, "localhost", 2000));
+
+        // Get the service provider.
+        var projectB = builder.AddProject<ProjectB>("b")
+            .WithEndpoints(projectA, "custom")
+            .WithReferenceEnvironment(flags);
+
+        // Call environment variable callbacks.
+        var config = await EnvironmentVariableEvaluator.GetEnvironmentVariablesAsync(projectB.Resource, DistributedApplicationOperation.Run, TestServiceProvider.Instance).DefaultTimeout();
+
+        switch (flags)
+        {
+            case ReferenceEnvironmentInjectionFlags.All:
+                Assert.Equal("https://localhost:2000", config["services__custom__mybinding__0"]);
+                Assert.Equal("https://localhost:2000", config["custom_MYBINDING"]);
+                break;
+            case ReferenceEnvironmentInjectionFlags.ConnectionProperties:
+                Assert.False(config.ContainsKey("custom_MYBINDING"));
+                Assert.False(config.ContainsKey("services__custom__mybinding__0"));
+                break;
+            case ReferenceEnvironmentInjectionFlags.ConnectionString:
+                Assert.False(config.ContainsKey("custom_MYBINDING"));
+                Assert.False(config.ContainsKey("services__custom__mybinding__0"));
+                break;
+            case ReferenceEnvironmentInjectionFlags.ServiceDiscovery:
+                Assert.False(config.ContainsKey("custom_MYBINDING"));
+                Assert.True(config.ContainsKey("services__custom__mybinding__0"));
+                break;
+            case ReferenceEnvironmentInjectionFlags.Endpoints:
+                Assert.True(config.ContainsKey("custom_MYBINDING"));
+                Assert.False(config.ContainsKey("services__custom__mybinding__0"));
+                break;
+            case ReferenceEnvironmentInjectionFlags.None:
+                Assert.False(config.ContainsKey("custom_MYBINDING"));
+                Assert.False(config.ContainsKey("services__custom__mybinding__0"));
+                break;
+        }
+    }
+
+    [Theory]
+    [InlineData(ReferenceEnvironmentInjectionFlags.All)]
+    [InlineData(ReferenceEnvironmentInjectionFlags.ConnectionProperties)]
+    [InlineData(ReferenceEnvironmentInjectionFlags.ConnectionString)]
+    [InlineData(ReferenceEnvironmentInjectionFlags.ServiceDiscovery)]
+    [InlineData(ReferenceEnvironmentInjectionFlags.Endpoints)]
+    [InlineData(ReferenceEnvironmentInjectionFlags.None)]
+    public async Task ContainerResourceWithEndpointRespectsCustomEnvironmentVariableNaming(ReferenceEnvironmentInjectionFlags flags)
+    {
+        using var builder = TestDistributedApplicationBuilder.Create();
+
+        // Create a binding and its matching annotation (simulating DCP behavior)
+        var container = builder.AddContainer("mycontainer", "myimage")
+                .WithHttpsEndpoint(1000, 2000, "mybinding")
+                .WithEndpoint("mybinding", e => e.AllocatedEndpoint = new AllocatedEndpoint(e, "localhost", 2000));
+
+        // Get the service provider.
+        var project = builder.AddProject<ProjectB>("b")
+            .WithEndpoints(container, "custom")
+            .WithReferenceEnvironment(flags);
+
+        // Call environment variable callbacks.
+        var config = await EnvironmentVariableEvaluator.GetEnvironmentVariablesAsync(project.Resource, DistributedApplicationOperation.Run, TestServiceProvider.Instance).DefaultTimeout();
+
+        switch (flags)
+        {
+            case ReferenceEnvironmentInjectionFlags.All:
+                Assert.Equal("https://localhost:2000", config["services__custom__mybinding__0"]);
+                Assert.Equal("https://localhost:2000", config["custom_MYBINDING"]);
+                break;
+            case ReferenceEnvironmentInjectionFlags.ConnectionProperties:
+                Assert.False(config.ContainsKey("custom_MYBINDING"));
+                Assert.False(config.ContainsKey("services__custom__mybinding__0"));
+                break;
+            case ReferenceEnvironmentInjectionFlags.ConnectionString:
+                Assert.False(config.ContainsKey("custom_MYBINDING"));
+                Assert.False(config.ContainsKey("services__custom__mybinding__0"));
+                break;
+            case ReferenceEnvironmentInjectionFlags.ServiceDiscovery:
+                Assert.False(config.ContainsKey("custom_MYBINDING"));
+                Assert.True(config.ContainsKey("services__custom__mybinding__0"));
+                break;
+            case ReferenceEnvironmentInjectionFlags.Endpoints:
+                Assert.True(config.ContainsKey("custom_MYBINDING"));
+                Assert.False(config.ContainsKey("services__custom__mybinding__0"));
+                break;
+            case ReferenceEnvironmentInjectionFlags.None:
+                Assert.False(config.ContainsKey("custom_MYBINDING"));
+                Assert.False(config.ContainsKey("services__custom__mybinding__0"));
+                break;
+        }
+    }
+
+    private sealed class ProjectA : IProjectMetadata
+    {
+        public string ProjectPath => "projectA";
+
+        public LaunchSettings LaunchSettings { get; } = new();
+    }
+
+    private sealed class ProjectB : IProjectMetadata
+    {
+        public string ProjectPath => "projectB";
+        public LaunchSettings LaunchSettings { get; } = new();
+    }
+}