lib/dns: simplify and add unit testing on zoneRecords

2 files changed, 134 insertions, 25 deletions

diff --git a/lib/dns/zone_records.go b/lib/dns/zone_records.go
index 63aeabbe..c7a28e1d 100644
--- a/lib/dns/zone_records.go
+++ b/lib/dns/zone_records.go
@@ -12,49 +12,46 @@ import "github.com/shuLhan/share/lib/reflect"
 //
 type zoneRecords map[string][]*ResourceRecord
 
-func (mr zoneRecords) add(rr *ResourceRecord) {
-	listRR := mr[rr.Name]
-
-	for x, rr2 := range listRR {
-		if rr.Type != rr2.Type {
-			continue
-		}
-		if rr.Class != rr2.Class {
-			continue
-		}
+//
+// add a ResourceRecord into the zone.
+//
+func (zr zoneRecords) add(rr *ResourceRecord) {
+	listRR := zr[rr.Name]
 
-		// Replace the RR if its type is SOA because only one SOA
-		// should exist per domain name.
-		if rr.Type == QueryTypeSOA {
+	// Replace the RR if its type is SOA because only one SOA
+	// should exist per domain name.
+	if rr.Type == QueryTypeSOA {
+		for x, in := range listRR {
+			if in.Type != QueryTypeSOA {
+				continue
+			}
 			listRR[x] = rr
 			return
 		}
-		break
 	}
 	listRR = append(listRR, rr)
-	mr[rr.Name] = listRR
+	zr[rr.Name] = listRR
 }
 
 //
-// remove a ResourceRecord from list of RR by its Name.
-// It will return true if the RR exist and removed, otherwise it will return
-// false.
+// remove a ResourceRecord from list by its Name and Value.
+// It will return true if the RR exist and removed.
 //
-func (mr zoneRecords) remove(rr *ResourceRecord) bool {
-	listRR := mr[rr.Name]
-	for x, rr2 := range listRR {
-		if rr.Type != rr2.Type {
+func (zr zoneRecords) remove(rr *ResourceRecord) bool {
+	listRR := zr[rr.Name]
+	for x, in := range listRR {
+		if in.Type != rr.Type {
 			continue
 		}
-		if rr.Class != rr2.Class {
+		if in.Class != rr.Class {
 			continue
 		}
-		if !reflect.IsEqual(rr.Value, rr2.Value) {
+		if !reflect.IsEqual(in.Value, rr.Value) {
 			continue
 		}
 		copy(listRR[x:], listRR[x+1:])
 		listRR[len(listRR)-1] = nil
-		mr[rr.Name] = listRR[:len(listRR)-1]
+		zr[rr.Name] = listRR[:len(listRR)-1]
 		return true
 	}
 	return false
diff --git a/lib/dns/zone_records_test.go b/lib/dns/zone_records_test.go
new file mode 100644
index 00000000..dc4ca22c
--- /dev/null
+++ b/lib/dns/zone_records_test.go
@@ -0,0 +1,112 @@
+// Copyright 2021, Shulhan <ms@kilabit.info>. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package dns
+
+import (
+	"testing"
+
+	"github.com/shuLhan/share/lib/test"
+)
+
+func testGenerateZoneRecords() (zoneRR zoneRecords, listRR []*ResourceRecord) {
+	zoneRR = zoneRecords{}
+
+	listRR = []*ResourceRecord{{
+		Name:  "test",
+		Type:  QueryTypeA,
+		Class: QueryClassIN,
+		Value: "127.0.0.1",
+		TTL:   1,
+	}, {
+		Name:  "test",
+		Type:  QueryTypeSOA,
+		Class: QueryClassIN,
+		Value: &RDataSOA{},
+		TTL:   2,
+	}, {
+		Name:  "test",
+		Type:  QueryTypeMX,
+		Class: QueryClassIN,
+		TTL:   3,
+	}, {
+		Name:  "test",
+		Type:  QueryTypeSOA,
+		Class: QueryClassIN,
+		TTL:   4,
+	}, {
+		Name:  "test",
+		Type:  QueryTypeA,
+		Class: QueryClassCH,
+		TTL:   5,
+	}}
+
+	for _, rr := range listRR {
+		zoneRR.add(rr)
+	}
+
+	return zoneRR, listRR
+}
+
+func TestZoneRecords_add(t *testing.T) {
+	gotZoneRR, listRR := testGenerateZoneRecords()
+
+	expZoneRR := zoneRecords{
+		"test": []*ResourceRecord{
+			listRR[0],
+			listRR[3],
+			listRR[2],
+			listRR[4],
+		},
+	}
+
+	test.Assert(t, "add", expZoneRR, gotZoneRR)
+}
+
+func TestZoneRecords_remove(t *testing.T) {
+	gotZoneRR, listRR := testGenerateZoneRecords()
+
+	cases := []struct {
+		rr           *ResourceRecord
+		expZoneRR    zoneRecords
+		expIsRemoved bool
+	}{{
+		// With different value.
+		rr: &ResourceRecord{
+			Name:  "test",
+			Type:  QueryTypeA,
+			Class: QueryClassIN,
+			Value: "127.0.0.2",
+		},
+		expZoneRR:    gotZoneRR,
+		expIsRemoved: false,
+	}, {
+		// With different Class.
+		rr: &ResourceRecord{
+			Name:  "test",
+			Type:  QueryTypeA,
+			Class: QueryClassCH,
+			Value: "127.0.0.1",
+		},
+		expZoneRR:    gotZoneRR,
+		expIsRemoved: false,
+	}, {
+		// With RR removed at the end.
+		rr: listRR[4],
+		expZoneRR: zoneRecords{
+			"test": []*ResourceRecord{
+				listRR[0],
+				listRR[3],
+				listRR[2],
+			},
+		},
+		expIsRemoved: true,
+	}}
+
+	for _, c := range cases {
+		gotIsRemoved := gotZoneRR.remove(c.rr)
+		test.Assert(t, "is removed", c.expIsRemoved, gotIsRemoved)
+		test.Assert(t, "after removed", c.expZoneRR, gotZoneRR)
+	}
+}