fluent-plugin-kusto 0.0.1.beta → 0.0.2.beta

data/test/plugin/test_e2e_kusto.rb

@@ -243,7 +243,7 @@ class KustoE2ETest < Test::Unit::TestCase
     kusto_query(".create table #{table_name} #{@columns}", :management)
   end
 
-  def wait_for_ingestion(query, expected_count, max_wait = 240, interval = 5)
+  def wait_for_ingestion(query, expected_count, max_wait = 480, interval = 5)
     waited = 0
     rows = []
 
@@ -339,15 +339,15 @@ class KustoE2ETest < Test::Unit::TestCase
       [time + 4, { 'id' => 2, 'name' => 'write_test_5' }]
     ]
 
-    @driver.run(default_tag: tag) do
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
       events.each do |t, r|
         @driver.feed(tag, t, r)
       end
-      sleep 5 # Wait for buffer flush
+      sleep 8 # Increased wait for buffer flush
     end
 
     query = "#{test_table} | extend r = parse_json(record) | where r.id == 2 and r.name startswith \"write_test_\""
-    rows = wait_for_ingestion(query, 5)
+    rows = wait_for_ingestion(query, 5, 600)  # Increased timeout to 10 minutes
 
     assert(rows.size >= 5, 'Not all events were ingested into Kusto by write')
   end
@@ -357,7 +357,8 @@ class KustoE2ETest < Test::Unit::TestCase
     configure_and_start_driver(
       table_name: test_table,
       buffered: true,
-      delayed: true
+      delayed: true,
+      deferred_commit_timeout: 120  # Increased timeout for chunk commit verification
     )
     setup_test_table(test_table)
 
@@ -371,15 +372,15 @@ class KustoE2ETest < Test::Unit::TestCase
       [time + 4, { 'id' => 3, 'name' => 'try_write_test_5' }]
     ]
 
-    @driver.run(default_tag: tag) do
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
       events.each do |t, r|
         @driver.feed(tag, t, r)
       end
-      sleep 5 # Wait for buffer flush
+      sleep 10 # Increased wait for buffer flush and delayed commit
     end
 
     query = "#{test_table} | extend r = parse_json(record) | where r.id == 3 and r.name startswith \"try_write_test_\""
-    rows = wait_for_ingestion(query, 5)
+    rows = wait_for_ingestion(query, 5, 600)  # Increased timeout to 10 minutes
 
     assert(rows.size >= 5, 'Not all events were ingested into Kusto by try_write')
 
@@ -387,7 +388,7 @@ class KustoE2ETest < Test::Unit::TestCase
     assert(chunk_id, 'chunk_id not found in ingested records')
 
     query_chunk = "#{test_table} | extend r = parse_json(record) | where r.chunk_id == '#{chunk_id}'"
-    chunk_rows = wait_for_ingestion(query_chunk, 5)
+    chunk_rows = wait_for_ingestion(query_chunk, 5, 600)  # Increased timeout to 10 minutes
 
     assert(chunk_rows.size >= 5, 'Not all chunk records were committed in Kusto by try_write')
   end
@@ -398,7 +399,8 @@ class KustoE2ETest < Test::Unit::TestCase
       table_name: test_table,
       buffered: true,
       delayed: true,
-      chunk_limit_size: '256'
+      chunk_limit_size: '256',
+      deferred_commit_timeout: 120  # Increased timeout for parallel chunk verification
     )
     setup_test_table(test_table)
 
@@ -409,15 +411,15 @@ class KustoE2ETest < Test::Unit::TestCase
       events << [time + i, { 'id' => 4, 'name' => "try_write_parallel_test_#{i + 1}" }]
     end
 
-    @driver.run(default_tag: tag) do
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
       events.each do |t, r|
         @driver.feed(tag, t, r)
       end
-      sleep 5 # Wait for buffer flush
+      sleep 10 # Increased wait for buffer flush and delayed commit
     end
 
     query = "#{test_table} | extend r = parse_json(record) | where r.id == 4 and r.name startswith \"try_write_parallel_test_\""
-    rows = wait_for_ingestion(query, 10)
+    rows = wait_for_ingestion(query, 10, 600)  # Increased timeout to 10 minutes
 
     assert(rows.size >= 10, 'Not all events were ingested into Kusto by try_write (parallel)')
 
@@ -428,7 +430,7 @@ class KustoE2ETest < Test::Unit::TestCase
     chunk_ids.each do |cid|
       expected_count = rows.count { |row| row[3]['chunk_id'] == cid }
       query_chunk = "#{test_table} | extend r = parse_json(record) | where r.chunk_id == '#{cid}'"
-      chunk_rows = wait_for_ingestion(query_chunk, expected_count)
+      chunk_rows = wait_for_ingestion(query_chunk, expected_count, 600)  # Increased timeout to 10 minutes
 
       assert(chunk_rows.size == expected_count,
              "Not all chunk records were committed in Kusto for chunk_id #{cid} (expected #{expected_count}, got #{chunk_rows.size})")
@@ -475,7 +477,7 @@ class KustoE2ETest < Test::Unit::TestCase
     tag = 'e2e.memory_buffered.immediate'
     events = generate_test_events(5, 2000, 'mem_imm')
 
-    @driver.run(default_tag: tag) do
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
       events.each do |time, record|
         @driver.feed(tag, time, record)
       end
@@ -483,7 +485,7 @@ class KustoE2ETest < Test::Unit::TestCase
     end
 
     query = "#{table_name} | extend r = parse_json(record) | where r.id >= 2000 and r.id <= 2004"
-    rows = wait_for_ingestion(query, 5)
+    rows = wait_for_ingestion(query, 5, 600)  # Increased timeout to 10 minutes
 
     assert(rows.size >= 5, "Expected 5 records, got #{rows.size} in memory buffered immediate flush")
   end
@@ -503,15 +505,15 @@ class KustoE2ETest < Test::Unit::TestCase
     tag = 'e2e.memory_buffered.interval'
     events = generate_test_events(7, 3000, 'mem_int')
 
-    @driver.run(default_tag: tag) do
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
       events.each do |time, record|
         @driver.feed(tag, time, record)
       end
-      sleep 8 # Wait longer than flush_interval
+      sleep 10 # Increased wait for buffer flush
     end
 
     query = "#{table_name} | extend r = parse_json(record) | where r.id >= 3000 and r.id <= 3006"
-    rows = wait_for_ingestion(query, 7)
+    rows = wait_for_ingestion(query, 7, 600)  # Increased timeout to 10 minutes
 
     assert(rows.size >= 7, "Expected 7 records, got #{rows.size} in memory buffered interval flush")
   end
@@ -543,15 +545,15 @@ class KustoE2ETest < Test::Unit::TestCase
       ]
     end
 
-    @driver.run(default_tag: tag) do
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
       events.each do |time, record|
         @driver.feed(tag, time, record)
       end
-      sleep 8
+      sleep 10  # Increased wait for buffer flush
     end
 
     query = "#{table_name} | extend r = parse_json(record) | where r.id >= 4000 and r.id <= 4009"
-    rows = wait_for_ingestion(query, 10)
+    rows = wait_for_ingestion(query, 10, 600)  # Increased timeout to 10 minutes
 
     assert(rows.size >= 10, "Expected 10 records, got #{rows.size} in chunk size limit test")
   end
@@ -563,23 +565,23 @@ class KustoE2ETest < Test::Unit::TestCase
       table_name: table_name,
       buffered: true,
       delayed: true,
-      flush_interval: '3s',
-      deferred_commit_timeout: 15
+      flush_interval: '5s',  # Longer flush interval to reduce buffer pressure
+      deferred_commit_timeout: 120  # Increased timeout for chunk commit verification
     )
     setup_test_table(table_name)
 
     tag = 'e2e.delayed_commit.sync'
     events = generate_test_events(4, 5000, 'delayed_sync')
 
-    @driver.run(default_tag: tag) do
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
       events.each do |time, record|
         @driver.feed(tag, time, record)
       end
-      sleep 8
+      sleep 15  # Increased wait for buffer flush and delayed commit
     end
 
     query = "#{table_name} | extend r = parse_json(record) | where r.id >= 5000 and r.id <= 5003"
-    rows = wait_for_ingestion(query, 4)
+    rows = wait_for_ingestion(query, 4, 600)  # Increased timeout to 10 minutes
 
     assert(rows.size >= 4, "Expected 4 records, got #{rows.size} in delayed commit sync mode")
 
@@ -595,26 +597,28 @@ class KustoE2ETest < Test::Unit::TestCase
       table_name: table_name,
       buffered: true,
       delayed: true,
-      chunk_limit_size: '300', # Small chunks to force multiple
-      flush_interval: '4s',
-      deferred_commit_timeout: 15
+      chunk_limit_size: '1k', # Reasonable size to force multiple chunks without overwhelming buffer
+      flush_interval: '6s',    # Longer flush interval to reduce buffer pressure
+      deferred_commit_timeout: 120,  # Increased timeout for chunk commit verification
+      flush_mode: 'interval'   # Ensure interval-based flushing
     )
     setup_test_table(table_name)
 
     tag = 'e2e.delayed_commit.multi_chunks'
-    events = generate_test_events(12, 6000, 'multi_chunk')
+    # Reduce number of events to prevent buffer overflow
+    events = generate_test_events(8, 6000, 'multi_chunk')
 
-    @driver.run(default_tag: tag) do
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
       events.each do |time, record|
         @driver.feed(tag, time, record)
       end
-      sleep 10
+      sleep 20  # Increased sleep time for buffer flush and delayed commit
     end
 
-    query = "#{table_name} | extend r = parse_json(record) | where r.id >= 6000 and r.id <= 6011"
-    rows = wait_for_ingestion(query, 12)
+    query = "#{table_name} | extend r = parse_json(record) | where r.id >= 6000 and r.id <= 6007"
+    rows = wait_for_ingestion(query, 8, 600)  # Increased timeout to 10 minutes
 
-    assert(rows.size >= 12, "Expected 12 records, got #{rows.size} in delayed commit multiple chunks")
+    assert(rows.size >= 8, "Expected 8 records, got #{rows.size} in delayed commit multiple chunks")
 
     # Verify multiple chunk_ids exist
     chunk_ids = rows.map { |row| row[3]['chunk_id'] if row[3] }.compact.uniq
@@ -638,15 +642,15 @@ class KustoE2ETest < Test::Unit::TestCase
     tag = 'e2e.file_buffer.persistent'
     events = generate_test_events(6, 20_000, 'file_buf')
 
-    @driver.run(default_tag: tag) do
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
       events.each do |time, record|
         @driver.feed(tag, time, record)
       end
-      sleep 8
+      sleep 10  # Increased wait for buffer flush
     end
 
     query = "#{table_name} | extend r = parse_json(record) | where r.id >= 20000 and r.id <= 20005"
-    rows = wait_for_ingestion(query, 6)
+    rows = wait_for_ingestion(query, 6, 600)  # Increased timeout to 10 minutes
 
     assert(rows.size >= 6, "Expected 6 records, got #{rows.size} in file buffer persistent storage test")
   end
@@ -666,18 +670,267 @@ class KustoE2ETest < Test::Unit::TestCase
     tag = 'e2e.buffered.compression'
     events = generate_test_events(10, 7000, 'compression')
 
-    @driver.run(default_tag: tag) do
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
       events.each do |time, record|
         @driver.feed(tag, time, record)
       end
-      sleep 8
+      sleep 10  # Increased wait for buffer flush
     end
 
     query = "#{table_name} | extend r = parse_json(record) | where r.id >= 7000 and r.id <= 7009"
-    rows = wait_for_ingestion(query, 10)
+    rows = wait_for_ingestion(query, 10, 600)  # Increased timeout to 10 minutes
 
     assert(rows.size >= 10, "Expected 10 records, got #{rows.size} in compression test")
   end
 
   # ESSENTIAL E2E BUFFERING TEST CASES - END
+
+  # INGESTION MAPPING REFERENCE TESTS - START
+  
+  # Test ingestion with mapping reference specified
+  test 'ingestion_with_mapping_reference' do
+    test_table = "FluentD_mapping_ref_#{Time.now.to_i}"
+    mapping_name = "test_mapping_#{Time.now.to_i}"
+    
+    # Create table and mapping
+    kusto_query(".drop table #{test_table} ifexists", :management)
+    kusto_query(".create table #{test_table} (tag:string, timestamp:datetime, record:dynamic)", :management)
+    kusto_query(<<~MAPPING_QUERY, :management)
+      .create table #{test_table} ingestion json mapping "#{mapping_name}" 
+      '[{"column":"tag","path":"$.tag"},{"column":"timestamp","path":"$.timestamp"},{"column":"record","path":"$.record"}]'
+    MAPPING_QUERY
+    
+    # Configure driver with mapping reference
+    config_options = {
+      table_name: test_table,
+      buffered: true,
+      delayed: true
+    }
+    
+    # Add ingestion_mapping_reference if specified
+    mapping_config = config_options[:ingestion_mapping_reference] ? "ingestion_mapping_reference #{config_options[:ingestion_mapping_reference]}" : ''
+    
+    @conf = <<-CONF
+      @type kusto
+      @log_level debug
+      buffered #{config_options[:buffered]}
+      delayed #{config_options[:delayed]}
+      endpoint #{@engine_url}
+      database_name #{@database}
+      table_name #{config_options[:table_name]}
+      compression_enabled true
+      ingestion_mapping_reference #{mapping_name}
+      #{@auth_lines}
+      <buffer>
+        @type memory
+        chunk_limit_size 8k
+        flush_interval 3s
+        flush_mode interval
+      </buffer>
+    CONF
+
+    @driver = Fluent::Test::Driver::Output.new(Fluent::Plugin::KustoOutput).configure(@conf)
+    @driver.instance.instance_variable_set(:@logger, @logger)
+    @driver.instance.start
+
+    tag = 'e2e.mapping_ref'
+    events = [
+      [Time.now.to_i, { 'id' => 8001, 'name' => 'mapping_test_1', 'type' => 'with_mapping' }],
+      [Time.now.to_i + 1, { 'id' => 8002, 'name' => 'mapping_test_2', 'type' => 'with_mapping' }],
+      [Time.now.to_i + 2, { 'id' => 8003, 'name' => 'mapping_test_3', 'type' => 'with_mapping' }]
+    ]
+
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
+      events.each do |time, record|
+        @driver.feed(tag, time, record)
+      end
+      sleep 10  # Increased wait for buffer flush
+    end
+
+    query = "#{test_table} | extend r = parse_json(record) | where r.id >= 8001 and r.id <= 8003"
+    rows = wait_for_ingestion(query, 3, 600)  # Increased timeout to 10 minutes
+
+    assert(rows.size >= 3, "Expected 3 records with mapping reference, got #{rows.size}")
+    
+    # Verify the mapping was used by checking data structure
+    found_with_mapping = false
+    rows.each do |row|
+      r = row[2] # record column should be dynamic
+      if r && r['id'] && r['id'] >= 8001 && r['id'] <= 8003
+        found_with_mapping = true
+        break
+      end
+    end
+    
+    assert(found_with_mapping, 'Expected records with mapping reference not found')
+    
+    # Clean up mapping
+    kusto_query(".drop table #{test_table} ingestion json mapping '#{mapping_name}'", :management)
+  end
+
+  # Test ingestion without mapping reference (default behavior)
+  test 'ingestion_without_mapping_reference' do
+    test_table = "FluentD_no_mapping_#{Time.now.to_i}"
+    
+    # Create table without specific mapping
+    kusto_query(".drop table #{test_table} ifexists", :management)
+    kusto_query(".create table #{test_table} (tag:string, timestamp:datetime, record:string)", :management)
+    
+    configure_and_start_driver(
+      table_name: test_table,
+      buffered: true,
+      delayed: false
+      # No ingestion_mapping_reference specified
+    )
+
+    tag = 'e2e.no_mapping'
+    events = [
+      [Time.now.to_i, { 'id' => 9001, 'name' => 'no_mapping_test_1', 'type' => 'default' }],
+      [Time.now.to_i + 1, { 'id' => 9002, 'name' => 'no_mapping_test_2', 'type' => 'default' }]
+    ]
+
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
+      events.each do |time, record|
+        @driver.feed(tag, time, record)
+      end
+      sleep 8  # Increased wait for buffer flush
+    end
+
+    query = "#{test_table} | extend r = parse_json(record) | where r.id >= 9001 and r.id <= 9002"
+    rows = wait_for_ingestion(query, 2, 600)  # Increased timeout to 10 minutes
+
+    assert(rows.size >= 2, "Expected 2 records without mapping reference, got #{rows.size}")
+    
+    # Verify default string serialization was used
+    found_default_format = false
+    rows.each do |row|
+      record_str = row[2] # record column should be string
+      if record_str.is_a?(String) && record_str.include?('"id":900')
+        found_default_format = true
+        break
+      end
+    end
+    
+    assert(found_default_format, 'Expected default JSON string format not found')
+  end
+
+  # Test ingestion mapping with delayed commit
+  test 'ingestion_mapping_with_delayed_commit' do
+    test_table = "FluentD_mapping_delayed_#{Time.now.to_i}"
+    mapping_name = "delayed_mapping_#{Time.now.to_i}"
+    
+    # Create table and mapping
+    kusto_query(".drop table #{test_table} ifexists", :management)
+    kusto_query(".create table #{test_table} (tag:string, timestamp:datetime, record:dynamic)", :management)
+    kusto_query(<<~MAPPING_QUERY, :management)
+      .create table #{test_table} ingestion json mapping "#{mapping_name}" 
+      '[{"column":"tag","path":"$.tag"},{"column":"timestamp","path":"$.timestamp"},{"column":"record","path":"$.record"}]'
+    MAPPING_QUERY
+    
+    # Configure with both mapping reference and delayed commit
+    @conf = <<-CONF
+      @type kusto
+      @log_level debug
+      buffered true
+      delayed true
+      endpoint #{@engine_url}
+      database_name #{@database}
+      table_name #{test_table}
+      compression_enabled true
+      ingestion_mapping_reference #{mapping_name}
+      deferred_commit_timeout 120
+      #{@auth_lines}
+      <buffer>
+        @type memory
+        chunk_limit_size 4k
+        flush_interval 5s
+        flush_mode interval
+      </buffer>
+    CONF
+
+    @driver = Fluent::Test::Driver::Output.new(Fluent::Plugin::KustoOutput).configure(@conf)
+    @driver.instance.instance_variable_set(:@logger, @logger)
+    @driver.instance.start
+
+    tag = 'e2e.mapping_delayed'
+    events = [
+      [Time.now.to_i, { 'id' => 10001, 'name' => 'delayed_mapping_1', 'type' => 'delayed_with_mapping' }],
+      [Time.now.to_i + 1, { 'id' => 10002, 'name' => 'delayed_mapping_2', 'type' => 'delayed_with_mapping' }],
+      [Time.now.to_i + 2, { 'id' => 10003, 'name' => 'delayed_mapping_3', 'type' => 'delayed_with_mapping' }]
+    ]
+
+    @driver.run(default_tag: tag, timeout: 300) do  # Increase driver timeout to 5 minutes
+      events.each do |time, record|
+        @driver.feed(tag, time, record)
+      end
+      sleep 12  # Increased wait for buffer flush and delayed commit
+    end
+
+    query = "#{test_table} | extend r = parse_json(record) | where r.id >= 10001 and r.id <= 10003"
+    rows = wait_for_ingestion(query, 3, 600)  # Increased timeout to 10 minutes
+
+    assert(rows.size >= 3, "Expected 3 records with mapping and delayed commit, got #{rows.size}")
+    
+    # Verify chunk_id exists (from delayed commit) and mapping was applied
+    chunk_ids = []
+    mapped_records = 0
+    
+    rows.each do |row|
+      r = row[2] # record column
+      if r && r['chunk_id'] # Should have chunk_id from delayed commit
+        chunk_ids << r['chunk_id']
+      end
+      if r && r['id'] && r['id'] >= 10001 && r['id'] <= 10003
+        mapped_records += 1
+      end
+    end
+    
+    assert(chunk_ids.uniq.size >= 1, 'Expected chunk_ids from delayed commit not found')
+    assert(mapped_records >= 3, 'Expected mapped records not found')
+    
+    # Clean up mapping
+    kusto_query(".drop table #{test_table} ingestion json mapping '#{mapping_name}'", :management)
+  end
+
+  # Test configuration validation for ingestion_mapping_reference
+  test 'ingestion_mapping_reference_configuration' do
+    test_table = "FluentD_config_test_#{Time.now.to_i}"
+    setup_test_table(test_table)
+    
+    # Test that plugin accepts ingestion_mapping_reference parameter
+    config_with_mapping = <<-CONF
+      @type kusto
+      buffered false
+      endpoint #{@engine_url}
+      database_name #{@database}
+      table_name #{test_table}
+      ingestion_mapping_reference test_mapping_name
+      #{@auth_lines}
+    CONF
+    
+    driver_with_mapping = nil
+    assert_nothing_raised('Configuration with ingestion_mapping_reference should be valid') do
+      driver_with_mapping = Fluent::Test::Driver::Output.new(Fluent::Plugin::KustoOutput).configure(config_with_mapping)
+    end
+    
+    # Verify the parameter is accessible
+    plugin_instance = driver_with_mapping.instance
+    assert_respond_to(plugin_instance, :ingestion_mapping_reference, 'Plugin should respond to ingestion_mapping_reference')
+    
+    # Test without mapping reference
+    config_without_mapping = <<-CONF
+      @type kusto
+      buffered false
+      endpoint #{@engine_url}
+      database_name #{@database}
+      table_name #{test_table}
+      #{@auth_lines}
+    CONF
+    
+    assert_nothing_raised('Configuration without ingestion_mapping_reference should be valid') do
+      Fluent::Test::Driver::Output.new(Fluent::Plugin::KustoOutput).configure(config_without_mapping)
+    end
+  end
+
+  # INGESTION MAPPING REFERENCE TESTS - END
 end

data/test/plugin/test_mi_tokenprovider.rb

@@ -0,0 +1,155 @@
+# frozen_string_literal: true
+
+require 'test/unit'
+require 'mocha/test_unit'
+require_relative '../../lib/fluent/plugin/auth/mi_tokenprovider'
+
+class DummyConfigForMI
+  attr_reader :kusto_endpoint, :managed_identity_client_id
+
+  def initialize(kusto_endpoint, managed_identity_client_id = nil)
+    @kusto_endpoint = kusto_endpoint
+    @managed_identity_client_id = managed_identity_client_id
+  end
+
+  def logger
+    require 'logger'
+    Logger.new($stdout)
+  end
+end
+
+class ManagedIdentityTokenProviderTest < Test::Unit::TestCase
+  def setup
+    @resource = 'https://kusto.kusto.windows.net'
+    @client_id = '074c3c54-29e2-4230-a81f-333868b8d6ca'
+  end
+
+  def test_initialize_with_user_managed_identity
+    config = DummyConfigForMI.new(@resource, @client_id)
+    provider = ManagedIdentityTokenProvider.new(config)
+    
+    # Verify instance variables are set correctly
+    assert_equal @resource, provider.instance_variable_get(:@resource)
+    assert_equal @client_id, provider.instance_variable_get(:@managed_identity_client_id)
+    assert_equal false, provider.instance_variable_get(:@use_system_assigned)
+    assert_equal true, provider.instance_variable_get(:@use_user_assigned)
+    assert_not_nil provider.instance_variable_get(:@token_acquire_url)
+  end
+
+  def test_initialize_with_system_managed_identity
+    config = DummyConfigForMI.new(@resource, 'SYSTEM')
+    provider = ManagedIdentityTokenProvider.new(config)
+    
+    # Verify instance variables are set correctly for system managed identity
+    assert_equal @resource, provider.instance_variable_get(:@resource)
+    assert_equal 'SYSTEM', provider.instance_variable_get(:@managed_identity_client_id)
+    assert_equal true, provider.instance_variable_get(:@use_system_assigned)
+    assert_equal false, provider.instance_variable_get(:@use_user_assigned)
+    assert_not_nil provider.instance_variable_get(:@token_acquire_url)
+  end
+
+  def test_initialize_with_empty_client_id
+    config = DummyConfigForMI.new(@resource, '')
+    provider = ManagedIdentityTokenProvider.new(config)
+    
+    # Verify instance variables are set correctly for empty client_id
+    assert_equal @resource, provider.instance_variable_get(:@resource)
+    assert_equal '', provider.instance_variable_get(:@managed_identity_client_id)
+    assert_equal false, provider.instance_variable_get(:@use_system_assigned)
+    assert_equal false, provider.instance_variable_get(:@use_user_assigned)
+  end
+
+  def test_token_acquire_url_formation_user_managed_identity
+    config = DummyConfigForMI.new(@resource, @client_id)
+    provider = ManagedIdentityTokenProvider.new(config)
+    
+    token_url = provider.instance_variable_get(:@token_acquire_url)
+    expected_base = 'http://169.254.169.254/metadata/identity/oauth2/token'
+    
+    assert token_url.include?(expected_base), "Token URL should contain base IMDS endpoint"
+    assert token_url.include?('resource=https%3A%2F%2Fkusto.kusto.windows.net'), "Token URL should contain encoded resource"
+    assert token_url.include?('api-version=2018-02-01'), "Token URL should contain API version"
+    assert token_url.include?("client_id=#{@client_id}"), "Token URL should contain client_id for user managed identity"
+  end
+
+  def test_token_acquire_url_formation_system_managed_identity
+    config = DummyConfigForMI.new(@resource, 'SYSTEM')
+    provider = ManagedIdentityTokenProvider.new(config)
+    
+    token_url = provider.instance_variable_get(:@token_acquire_url)
+    expected_base = 'http://169.254.169.254/metadata/identity/oauth2/token'
+    
+    assert token_url.include?(expected_base), "Token URL should contain base IMDS endpoint"
+    assert token_url.include?('resource=https%3A%2F%2Fkusto.kusto.windows.net'), "Token URL should contain encoded resource"
+    assert token_url.include?('api-version=2018-02-01'), "Token URL should contain API version"
+    assert !token_url.include?('client_id='), "Token URL should NOT contain client_id for system managed identity"
+  end
+
+  def test_fetch_token_success
+    config = DummyConfigForMI.new(@resource, @client_id)
+    provider = ManagedIdentityTokenProvider.new(config)
+    
+    # Mock successful HTTP response
+    mock_response = {
+      'access_token' => 'fake-access-token',
+      'expires_in' => 3600
+    }
+    
+    provider.stubs(:post_token_request).returns(mock_response)
+    
+    result = provider.send(:fetch_token)
+    assert_equal 'fake-access-token', result[:access_token]
+    assert_equal 3600, result[:expires_in]
+  end
+
+  def test_post_token_request_retries_on_failure
+    config = DummyConfigForMI.new(@resource, @client_id)
+    provider = ManagedIdentityTokenProvider.new(config)
+    
+    # Mock HTTP failure followed by success
+    mock_http = mock
+    mock_response_fail = mock
+    mock_response_fail.stubs(:code).returns(500)
+    mock_response_fail.stubs(:body).returns('Internal Server Error')
+    
+    mock_response_success = mock
+    mock_response_success.stubs(:code).returns(200)
+    mock_response_success.stubs(:body).returns('{"access_token":"fake-token","expires_in":3600}')
+    
+    mock_request = mock
+    Net::HTTP::Get.stubs(:new).returns(mock_request)
+    
+    mock_http.expects(:request).twice.returns(mock_response_fail, mock_response_success)
+    Net::HTTP.stubs(:new).returns(mock_http)
+    
+    # Stub sleep to speed up test
+    provider.stubs(:sleep)
+    
+    result = provider.send(:post_token_request)
+    assert_equal 'fake-token', result['access_token']
+  end
+
+  def test_post_token_request_raises_after_max_retries
+    config = DummyConfigForMI.new(@resource, @client_id)
+    provider = ManagedIdentityTokenProvider.new(config)
+    
+    # Mock HTTP failure for all attempts
+    mock_http = mock
+    mock_response_fail = mock
+    mock_response_fail.stubs(:code).returns(500)
+    mock_response_fail.stubs(:body).returns('Internal Server Error')
+    
+    mock_request = mock
+    Net::HTTP::Get.stubs(:new).returns(mock_request)
+    
+    mock_http.stubs(:request).returns(mock_response_fail)
+    Net::HTTP.stubs(:new).returns(mock_http)
+    
+    # Stub sleep to speed up test
+    provider.stubs(:sleep)
+    
+    assert_raise(RuntimeError, 'Failed to get managed identity token after 2 attempts.') do
+      provider.send(:post_token_request)
+    end
+  end
+end