require 'test_helper'
require 'net/dns/resolver'

class TestDNSResolver < Minitest::Test
  def setup
    @resolver = Net::DNS::Resolver.new(
        nameservers: ['8.8.8.8'],
        port: 53,
        source_address: IPAddr.new('0.0.0.0'),
        tcp_timeout: 5,
        udp_timeout: 5
      )

    # Sample DNS packet data
    @packet_data = "\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\x03www\x06google\x03com\x00\x00\x01\x00\x01"
  end

  # Tests for query_tcp
  def test_query_tcp_successful_response
    mock_socket = Minitest::Mock.new

    Socket.stub(:new, mock_socket) do
      mock_socket.expect(:bind, nil, [String])
      mock_socket.expect(:connect, nil, [String])
      mock_socket.expect(:write, @packet_data.size + 2, [String])
      mock_socket.expect(:recv, "\x00\x20", [2]) # Length = 32 bytes
      mock_socket.expect(:recvfrom, ["A" * 32, nil], [32])
      mock_socket.expect(:close, nil)

      Socket.stub(:pack_sockaddr_in, "mock_sockaddr") do
        result = @resolver.send(:query_tcp, @packet, @packet_data)

        refute_nil result
        assert_equal 2, result.size
        assert_equal "A" * 32, result[0] # Response data
        assert_equal ["", 53, "8.8.8.8", "8.8.8.8"], result[1] # From info
      end
    end

    mock_socket.verify
  end

  def test_query_tcp_zero_length_response
    mock_socket = Minitest::Mock.new

    Socket.stub(:new, mock_socket) do
      mock_socket.expect(:bind, nil, [String])
      mock_socket.expect(:connect, nil, [String])
      mock_socket.expect(:write, @packet_data.size + 2, [String])
      mock_socket.expect(:recv, "\x00\x00", [2]) # Length = 0 bytes
      mock_socket.expect(:close, nil)

      Socket.stub(:pack_sockaddr_in, "mock_sockaddr") do
        result = @resolver.send(:query_tcp, @packet, @packet_data)
        assert_nil result # Should return nil for zero-length response
      end
    end

    mock_socket.verify
  end

  def test_query_tcp_bogus_response
    mock_socket = Minitest::Mock.new

    Socket.stub(:new, mock_socket) do
      mock_socket.expect(:bind, nil, [String])
      mock_socket.expect(:connect, nil, [String])
      mock_socket.expect(:write, @packet_data.size + 2, [String])
      mock_socket.expect(:recv, "0", [2]) # Length = 0 bytes
      mock_socket.expect(:close, nil)

      Socket.stub(:pack_sockaddr_in, "mock_sockaddr") do
        result = @resolver.send(:query_tcp, @packet, @packet_data)
        assert_nil result # Should return nil for zero-length response
      end
    end

    mock_socket.verify
  end
end
