From 8c5a65a915bde90209f2f128e3398a23195b82d4 Mon Sep 17 00:00:00 2001
From: Adam Shaw <atshaw@amazon.com>
Date: Tue, 10 Oct 2023 22:30:22 -0700
Subject: [PATCH] Addressing PR

---
 .../internal/translator/segment.go            | 26 +++++++++----------
 .../internal/translator/segment_test.go       | 14 +++-------
 2 files changed, 16 insertions(+), 24 deletions(-)

diff --git a/exporter/awsxrayexporter/internal/translator/segment.go b/exporter/awsxrayexporter/internal/translator/segment.go
index 5bae164a639d..21470e0e0dbd 100644
--- a/exporter/awsxrayexporter/internal/translator/segment.go
+++ b/exporter/awsxrayexporter/internal/translator/segment.go
@@ -116,11 +116,11 @@ func isLocalRoot(span ptrace.Span) bool {
 	return false
 }
 
-func addNamespaceToSpansWithRemoteService(span ptrace.Span, segment *awsxray.Segment) {
-	// We need to at least cover the
+func addNamespaceToSubsegmentWithRemoteService(span ptrace.Span, segment *awsxray.Segment) {
 	if (span.Kind() == ptrace.SpanKindClient ||
 		span.Kind() == ptrace.SpanKindConsumer ||
 		span.Kind() == ptrace.SpanKindProducer) &&
+		segment.Type != nil &&
 		segment.Namespace == nil {
 		if _, ok := span.Attributes().Get(awsRemoteService); ok {
 			segment.Namespace = awsxray.String("remote")
@@ -143,13 +143,19 @@ func MakeDependencySubsegmentForLocalRootDependencySpan(span ptrace.Span, resour
 	// Make this a subsegment
 	dependencySubsegment.Type = awsxray.String("subsegment")
 
-	addNamespaceToSpansWithRemoteService(span, dependencySubsegment)
+	if dependencySubsegment.Namespace == nil {
+		dependencySubsegment.Namespace = awsxray.String("remote")
+	}
 
 	// Remove span links from consumer spans
 	if span.Kind() == ptrace.SpanKindConsumer {
 		dependencySubsegment.Links = nil
 	}
 
+	myAwsRemoteService, _ := span.Attributes().Get(awsRemoteService)
+
+	dependencySubsegment.Name = awsxray.String(myAwsRemoteService.Str())
+
 	return dependencySubsegment, err
 }
 
@@ -218,14 +224,8 @@ func MakeServiceSegmentForLocalRootSpanWithoutDependency(span ptrace.Span, resou
 		return nil, err
 	}
 
-	// Make internal spans a segment
-	// And Ensure consumer process spans are not made a segment
-	_, hasAwsRemoteService := span.Attributes().Get(awsRemoteService)
-
-	if hasAwsRemoteService {
-		segment.Type = nil
-		segment.Namespace = nil
-	}
+	segment.Type = nil
+	segment.Namespace = nil
 
 	return []*awsxray.Segment{segment}, err
 }
@@ -237,7 +237,7 @@ func MakeNonLocalRootSegment(span ptrace.Span, resource pcommon.Resource, indexe
 		return nil, err
 	}
 
-	addNamespaceToSpansWithRemoteService(span, segment)
+	addNamespaceToSubsegmentWithRemoteService(span, segment)
 
 	return []*awsxray.Segment{segment}, nil
 }
@@ -356,7 +356,7 @@ func MakeSegment(span ptrace.Span, resource pcommon.Resource, indexedAttrs []str
 		}
 	}
 
-	if span.Kind() == ptrace.SpanKindClient || span.Kind() == ptrace.SpanKindProducer || span.Kind() == ptrace.SpanKindConsumer {
+	if span.Kind() == ptrace.SpanKindClient || span.Kind() == ptrace.SpanKindProducer {
 		if remoteServiceName, ok := attributes.Get(awsRemoteService); ok {
 			name = remoteServiceName.Str()
 		}
diff --git a/exporter/awsxrayexporter/internal/translator/segment_test.go b/exporter/awsxrayexporter/internal/translator/segment_test.go
index c2144d049531..a9ca076c3659 100644
--- a/exporter/awsxrayexporter/internal/translator/segment_test.go
+++ b/exporter/awsxrayexporter/internal/translator/segment_test.go
@@ -1024,20 +1024,12 @@ func TestConsumerSpanWithAwsRemoteServiceName(t *testing.T) {
 	spanName := "ABC.payment"
 	parentSpanID := newSegmentID()
 	user := "testingT"
-	attributes := make(map[string]interface{})
-	attributes[conventions.AttributeHTTPMethod] = "POST"
-	attributes[conventions.AttributeHTTPScheme] = "https"
-	attributes[conventions.AttributeHTTPHost] = "payment.amazonaws.com"
-	attributes[conventions.AttributeHTTPTarget] = "/"
-	attributes[conventions.AttributeRPCService] = "ABC"
+	attributes := getBasicAttributes()
 	attributes[awsRemoteService] = "ConsumerService"
 
 	resource := constructDefaultResource()
 	span := constructConsumerSpan(parentSpanID, spanName, 0, "Ok", attributes)
 
-	segment, _ := MakeSegment(span, resource, nil, false, nil, false)
-	assert.Equal(t, "ConsumerService", *segment.Name)
-
 	jsonStr, err := MakeSegmentDocumentString(span, resource, nil, false, nil, false)
 
 	assert.NotNil(t, jsonStr)
@@ -1213,7 +1205,7 @@ func TestLocalRootConsumer(t *testing.T) {
 	assert.Equal(t, segments[0].EndTime, segments[1].EndTime)
 }
 
-func TestLocalRootConsumerProcess(t *testing.T) {
+func TestNonLocalRootConsumerProcess(t *testing.T) {
 	spanName := "destination operation"
 	resource := getBasicResource()
 	parentSpanID := newSegmentID()
@@ -1472,7 +1464,7 @@ func TestNotLocalRootConsumer(t *testing.T) {
 	// Validate segment
 	assert.Equal(t, "subsegment", *segments[0].Type)
 	assert.Equal(t, "remote", *segments[0].Namespace)
-	assert.Equal(t, "myRemoteService", *segments[0].Name)
+	assert.Equal(t, "MyService", *segments[0].Name)
 }
 
 func TestNotLocalRootClient(t *testing.T) {