package dns

import (
	"testing"
	"time"

	"github.com/miekg/dns"
)

func TestCacheSetGet(t *testing.T) {
	c := NewCache()

	msg := new(dns.Msg)
	msg.SetQuestion("example.com.", dns.TypeA)
	msg.Answer = append(msg.Answer, &dns.A{
		Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60},
		A:   []byte{1, 2, 3, 4},
	})

	c.Set("example.com.", dns.TypeA, dns.ClassINET, msg, 5*time.Second)

	cached := c.Get("example.com.", dns.TypeA, dns.ClassINET)
	if cached == nil {
		t.Fatal("expected cached response")
	}
	if len(cached.Answer) != 1 {
		t.Fatalf("got %d answers, want 1", len(cached.Answer))
	}
}

func TestCacheMiss(t *testing.T) {
	c := NewCache()

	cached := c.Get("example.com.", dns.TypeA, dns.ClassINET)
	if cached != nil {
		t.Fatal("expected nil for cache miss")
	}
}

func TestCacheExpiry(t *testing.T) {
	c := NewCache()

	msg := new(dns.Msg)
	msg.SetQuestion("example.com.", dns.TypeA)

	c.Set("example.com.", dns.TypeA, dns.ClassINET, msg, 1*time.Millisecond)
	time.Sleep(2 * time.Millisecond)

	cached := c.Get("example.com.", dns.TypeA, dns.ClassINET)
	if cached != nil {
		t.Fatal("expected nil for expired entry")
	}
}

func TestCacheDifferentTypes(t *testing.T) {
	c := NewCache()

	msgA := new(dns.Msg)
	msgA.SetQuestion("example.com.", dns.TypeA)
	c.Set("example.com.", dns.TypeA, dns.ClassINET, msgA, 5*time.Second)

	msgAAAA := new(dns.Msg)
	msgAAAA.SetQuestion("example.com.", dns.TypeAAAA)
	c.Set("example.com.", dns.TypeAAAA, dns.ClassINET, msgAAAA, 5*time.Second)

	cachedA := c.Get("example.com.", dns.TypeA, dns.ClassINET)
	if cachedA == nil {
		t.Fatal("expected cached A response")
	}

	cachedAAAA := c.Get("example.com.", dns.TypeAAAA, dns.ClassINET)
	if cachedAAAA == nil {
		t.Fatal("expected cached AAAA response")
	}

	// Different type should not match.
	cachedMX := c.Get("example.com.", dns.TypeMX, dns.ClassINET)
	if cachedMX != nil {
		t.Fatal("expected nil for uncached type")
	}
}