Add tests

Author: JeffreyCA

cli/azd/internal/cmd/errors_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/azure/azure-dev/cli/azd/internal/tracing/fields"
 	"github.com/azure/azure-dev/cli/azd/pkg/auth"
 	"github.com/azure/azure-dev/cli/azd/pkg/azapi"
+	"github.com/azure/azure-dev/cli/azd/pkg/azdext"
 	"github.com/azure/azure-dev/cli/azd/pkg/exec"
 	"github.com/azure/azure-dev/cli/azd/test/mocks/mocktracing"
 	"github.com/stretchr/testify/require"
@@ -148,6 +149,23 @@ func Test_MapError(t *testing.T) {
 				fields.ErrorKey(fields.ServiceCorrelationId.Key).String("12345"),
 			},
 		},
+		{
+			name: "WithExtensionResponseError",
+			err: &azdext.ExtensionResponseError{
+				Message:     "Rate limit exceeded",
+				Details:     "Too many requests",
+				ErrorCode:   "RateLimitExceeded",
+				StatusCode:  429,
+				ServiceName: "openai.azure.com",
+			},
+			wantErrReason: "ext.service.openai.429",
+			wantErrDetails: []attribute.KeyValue{
+				fields.ErrorKey(fields.ServiceName.Key).String("openai"),
+				fields.ErrorKey(fields.ServiceHost.Key).String("openai.azure.com"),
+				fields.ErrorKey(fields.ServiceStatusCode.Key).Int(429),
+				fields.ErrorKey(fields.ServiceErrorCode.Key).String("RateLimitExceeded"),
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

cli/azd/pkg/azdext/extension_error_test.go

@@ -0,0 +1,118 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azdext
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestExtensionError_RoundTrip(t *testing.T) {
+	tests := []struct {
+		name                    string
+		inputErr                error
+		wantInfo                errorInfo
+		useErrorStringAsMessage bool
+	}{
+		{
+			name:     "NilError",
+			inputErr: nil,
+			wantInfo: errorInfo{},
+		},
+		{
+			name:     "SimpleError",
+			inputErr: errors.New("simple error"),
+			wantInfo: errorInfo{
+				message: "simple error",
+			},
+		},
+		{
+			name: "ExtensionResponseError",
+			inputErr: &ExtensionResponseError{
+				Message:     "Rate limit exceeded",
+				Details:     "Too many requests",
+				ErrorCode:   "RateLimitExceeded",
+				StatusCode:  429,
+				ServiceName: "openai.azure.com",
+			},
+			wantInfo: errorInfo{
+				message:    "Rate limit exceeded",
+				details:    "Too many requests",
+				errorCode:  "RateLimitExceeded",
+				statusCode: 429,
+				service:    "openai.azure.com",
+			},
+		},
+		{
+			name: "AzCoreResponseError",
+			inputErr: &azcore.ResponseError{
+				ErrorCode:  "ResourceNotFound",
+				StatusCode: 404,
+			},
+			wantInfo: errorInfo{
+				errorCode:  "ResourceNotFound",
+				statusCode: 404,
+			},
+			useErrorStringAsMessage: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.useErrorStringAsMessage {
+				tt.wantInfo.message = tt.inputErr.Error()
+			}
+
+			// Helper to verify the proto message content
+			type protoMessage interface {
+				GetMessage() string
+				GetDetails() string
+				GetErrorCode() string
+				GetStatusCode() int32
+				GetServiceName() string
+			}
+			verifyProto := func(t *testing.T, msg protoMessage) {
+				assert.NotNil(t, msg)
+				assert.Equal(t, tt.wantInfo.message, msg.GetMessage())
+				assert.Equal(t, tt.wantInfo.details, msg.GetDetails())
+				assert.Equal(t, tt.wantInfo.errorCode, msg.GetErrorCode())
+				assert.Equal(t, tt.wantInfo.statusCode, msg.GetStatusCode())
+				assert.Equal(t, tt.wantInfo.service, msg.GetServiceName())
+			}
+
+			// Helper to verify the unwrapped error
+			verifyUnwrapped := func(t *testing.T, err error) {
+				var extErr *ExtensionResponseError
+				if assert.ErrorAs(t, err, &extErr) {
+					assert.Equal(t, tt.wantInfo.message, extErr.Message)
+					assert.Equal(t, tt.wantInfo.details, extErr.Details)
+					assert.Equal(t, tt.wantInfo.errorCode, extErr.ErrorCode)
+					assert.Equal(t, int(tt.wantInfo.statusCode), extErr.StatusCode)
+					assert.Equal(t, tt.wantInfo.service, extErr.ServiceName)
+				}
+			}
+
+			// Test ServiceTarget wrapping
+			stMsg := WrapErrorForServiceTarget(tt.inputErr)
+			if tt.inputErr == nil {
+				assert.Nil(t, stMsg)
+			} else {
+				verifyProto(t, stMsg)
+				verifyUnwrapped(t, UnwrapErrorFromServiceTarget(stMsg))
+			}
+
+			// Test FrameworkService wrapping
+			fsMsg := WrapErrorForFrameworkService(tt.inputErr)
+			if tt.inputErr == nil {
+				assert.Nil(t, fsMsg)
+			} else {
+				verifyProto(t, fsMsg)
+				verifyUnwrapped(t, UnwrapErrorFromFrameworkService(fsMsg))
+			}
+		})
+	}
+}

cli/azd/pkg/grpcbroker/message_broker_test.go

@@ -24,6 +24,7 @@ type TestMessage struct {
 	InnerMsg     any
 	IsProgress   bool
 	ProgressText string
+	TraceParent  string
 }
 
 // Test request/response types for handler testing
@@ -164,6 +165,14 @@ func (e *SimpleMessageEnvelope) CreateProgressMessage(requestId string, message
 	}
 }
 
+func (e *SimpleMessageEnvelope) GetTraceParent(msg *TestMessage) string {
+	return msg.TraceParent
+}
+
+func (e *SimpleMessageEnvelope) SetTraceParent(msg *TestMessage, traceParent string) {
+	msg.TraceParent = traceParent
+}
+
 // TestOn_RegistersHandler tests that handlers are registered correctly
 func TestOn_RegistersHandler(t *testing.T) {
 	sim := NewSimulatedBidiStream()