From e84b23d2fc13f7d7a2e5965066be871f5ede314e Mon Sep 17 00:00:00 2001
From: Sam Handler <sam.handler@shopify.com>
Date: Thu, 15 Feb 2024 15:53:05 -0700
Subject: [PATCH] fixup tests

---
 .../trilogy/instrumentation_test.rb           | 49 ++++++++++++++++---
 1 file changed, 42 insertions(+), 7 deletions(-)

diff --git a/instrumentation/trilogy/test/opentelemetry/instrumentation/trilogy/instrumentation_test.rb b/instrumentation/trilogy/test/opentelemetry/instrumentation/trilogy/instrumentation_test.rb
index a7f949f3d..1524802aa 100644
--- a/instrumentation/trilogy/test/opentelemetry/instrumentation/trilogy/instrumentation_test.rb
+++ b/instrumentation/trilogy/test/opentelemetry/instrumentation/trilogy/instrumentation_test.rb
@@ -32,7 +32,7 @@
   let(:port) { ENV.fetch('TEST_MYSQL_PORT', '3306').to_i }
   let(:database) { ENV.fetch('TEST_MYSQL_DB', 'mysql') }
   let(:username) { ENV.fetch('TEST_MYSQL_USER', 'root') }
-  let(:password) { ENV.fetch('TEST_MYSQL_PASSWORD', 'root') }
+  let(:password) { ENV.fetch('TEST_MYSQL_PASSWORD', '') }
 
   before do
     exporter.reset
@@ -375,15 +375,50 @@
     describe 'when propagator is set to vitess' do
       let(:config) { { propagator: 'vitess' } }
 
-      it 'does inject context' do
-        sql = 'SELECT * from users where users.id = 1 and users.email = "test@test.com"'.freeze # rubocop:disable Style/RedundantFreeze
-        propagator = double('propagator')
-        allow(propagator).to receive(:inject)
-        allow(client).to receive(:propagator).and_return(propagator)
+      it 'does inject context on frozen strings' do
+        sql = 'SELECT * from users where users.id = 1 and users.email = "test@test.com"'
+        propagator = OpenTelemetry::Instrumentation::Trilogy::Instrumentation.instance.propagator
+
+        arg_cache = {} # maintain handles to args
+        allow(client).to receive(:query).and_wrap_original do |m, *args|
+          arg_cache[:query_input] = args[0]
+          assert(args[0].frozen?)
+          m.call(args[0])
+        end
+
+        allow(propagator).to receive(:inject).and_wrap_original do |m, *args|
+          arg_cache[:inject_input] = args[0]
+          refute(args[0].frozen?)
+          assert_match(sql, args[0])
+          m.call(args[0], context: args[1][:context])
+        end
+
         expect do
           client.query(sql)
         end.must_raise Trilogy::Error
-        expect(propagator).to have_received(:inject).with(sql, context: instance_of(OpenTelemetry::Context)).once
+
+        # arg_cache[:inject_input] _was_ a mutable string, so it has the context injected
+        encoded = Base64.strict_encode64("{\"uber-trace-id\":\"#{span.hex_trace_id}:#{span.hex_span_id}:0:1\"}")
+        assert_equal(arg_cache[:inject_input], "/*VT_SPAN_CONTEXT=#{encoded}*/#{sql}")
+
+        # arg_cache[:inject_input] is now frozen
+        assert(arg_cache[:inject_input].frozen?)
+      end
+
+      it 'does inject context on unfrozen strings' do
+        # inbound SQL is not frozen (string prefixed with +)
+        sql = +'SELECT * from users where users.id = 1 and users.email = "test@test.com"'
+
+        # dup sql for comparison purposes, since propagator  mutates it
+        cached_sql = sql.dup
+
+        expect do
+          client.query(sql)
+        end.must_raise Trilogy::Error
+
+        encoded = Base64.strict_encode64("{\"uber-trace-id\":\"#{span.hex_trace_id}:#{span.hex_span_id}:0:1\"}")
+        assert_equal(sql, "/*VT_SPAN_CONTEXT=#{encoded}*/#{cached_sql}")
+        refute(sql.frozen?)
       end
     end