dd-trace-java/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbUtilTest.groovy at 2645dc6f73140ef11d65173fbae4a077090315a5 · DataDog/dd-trace-java · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink
import datadog.trace.bootstrap.instrumentation.api.SpanPointerUtils
import datadog.trace.instrumentation.aws.v2.dynamodb.DynamoDbUtil
import org.junit.jupiter.api.Test
import software.amazon.awssdk.core.SdkBytes
import software.amazon.awssdk.services.dynamodb.model.AttributeValue

class DynamoDbUtilTest {
  static createMockSpan() {
    def links = []

    def mockSpan = [
      addLink: { AgentSpanLink link ->
        links.add(link)
      }
    ] as AgentSpan

    return [span: mockSpan, links: links]
  }

  @Test
  void testAddSpanPointerWithNullKeys() {
    def mockData = createMockSpan()
    def mockSpan = mockData.span
    def links = mockData.links

    DynamoDbUtil.addSpanPointer(mockSpan, "table", null)

    assert links.isEmpty()
  }

  @Test
  void testAddSpanPointerWithEmptyKeys() {
    def mockData = createMockSpan()
    def mockSpan = mockData.span
    def links = mockData.links

    DynamoDbUtil.addSpanPointer(mockSpan, "table", [:])

    assert links.isEmpty()
  }

  @Test
  void testAddSpanPointerWithNullTableName() {
    def mockData = createMockSpan()
    def mockSpan = mockData.span
    def links = mockData.links

    def keys = [
      "id": AttributeValue.builder().s("12345").build()
    ]

    DynamoDbUtil.addSpanPointer(mockSpan, null, keys)

    assert links.isEmpty()
  }

  @Test
  void testAddSpanPointerWithSingleStringKey() {
    def mockData = createMockSpan()
    def mockSpan = mockData.span
    def links = mockData.links

    def keys = [
      "id": AttributeValue.builder().s("12345").build()
    ]

    DynamoDbUtil.addSpanPointer(mockSpan, "my-table", keys)

    assert links.size() == 1
    def link = links[0]
    assert link.attributes().asMap().get("ptr.kind") == SpanPointerUtils.DYNAMODB_PTR_KIND
    assert link.attributes().asMap().get("ptr.dir") == SpanPointerUtils.DOWN_DIRECTION
    assert link.attributes().asMap().get("link.kind") == SpanPointerUtils.LINK_KIND
    assert link.attributes().asMap().get("ptr.hash") != null
  }

  @Test
  void testAddSpanPointerWithSingleNumberKey() {
    def mockData = createMockSpan()
    def mockSpan = mockData.span
    def links = mockData.links

    def keys = [
      "count": AttributeValue.builder().n("42").build()
    ]

    DynamoDbUtil.addSpanPointer(mockSpan, "my-table", keys)

    assert links.size() == 1
    def link = links[0]
    assert link.attributes().asMap().get("ptr.kind") == SpanPointerUtils.DYNAMODB_PTR_KIND
  }

  @Test
  void testAddSpanPointerWithSingleBinaryKey() {
    def mockData = createMockSpan()
    def mockSpan = mockData.span
    def links = mockData.links

    def binaryData = "binary-data".getBytes()
    def keys = [
      "data": AttributeValue.builder().b(SdkBytes.fromByteArray(binaryData)).build()
    ]

    DynamoDbUtil.addSpanPointer(mockSpan, "my-table", keys)

    assert links.size() == 1
    def link = links[0]
    assert link.attributes().asMap().get("ptr.kind") == SpanPointerUtils.DYNAMODB_PTR_KIND
  }

  @Test
  void testAddSpanPointerWithTwoKeys() {
    def mockData = createMockSpan()
    def mockSpan = mockData.span
    def links = mockData.links

    def keys = [
      "id": AttributeValue.builder().s("12345").build(),
      "name": AttributeValue.builder().s("item-name").build()
    ]

    DynamoDbUtil.addSpanPointer(mockSpan, "my-table", keys)

    assert links.size() == 1
    def link = links[0]
    assert link.attributes().asMap().get("ptr.kind") == SpanPointerUtils.DYNAMODB_PTR_KIND
    assert link.attributes().asMap().get("ptr.hash") != null
  }

  @Test
  void testAddSpanPointerWithTwoKeysSortsAlphabetically() {
    def mockData = createMockSpan()
    def mockSpan1 = mockData.span
    def links1 = mockData.links

    // Keys in order bKey, aKey
    def keys1 = [
      "bKey": AttributeValue.builder().s("abc").build(),
      "aKey": AttributeValue.builder().s("zxy").build()
    ]

    DynamoDbUtil.addSpanPointer(mockSpan1, "my-table", keys1)

    // Reverse order: aKey, bKey — should produce the same hash
    def mockData2 = createMockSpan()
    def mockSpan2 = mockData2.span
    def links2 = mockData2.links

    def keys2 = [
      "aKey": AttributeValue.builder().s("zxy").build(),
      "bKey": AttributeValue.builder().s("abc").build()
    ]

    DynamoDbUtil.addSpanPointer(mockSpan2, "my-table", keys2)

    assert links1.size() == 1
    assert links2.size() == 1
    // Both should produce the same hash regardless of input order
    assert links1[0].attributes().asMap().get("ptr.hash") == links2[0].attributes().asMap().get("ptr.hash")
  }
}