1 files changed, 18 insertions, 41 deletions

diff --git a/dhcpv4/dhcpv4_test.go b/dhcpv4/dhcpv4_test.go
index e5b8ad6..fb0ef70 100644
--- a/dhcpv4/dhcpv4_test.go
+++ b/dhcpv4/dhcpv4_test.go
@@ -180,64 +180,42 @@ func TestGetOption(t *testing.T) {
 	}
 
 	hostnameOpt := &OptionGeneric{OptionCode: OptionHostName, Data: []byte("darkstar")}
-	bootFileOpt1 := &OptBootfileName{"boot.img"}
-	bootFileOpt2 := &OptBootfileName{"boot2.img"}
-	d.AddOption(hostnameOpt)
-	d.AddOption(&OptBootfileName{"boot.img"})
-	d.AddOption(&OptBootfileName{"boot2.img"})
-
-	require.Equal(t, d.GetOption(OptionHostName), []Option{hostnameOpt})
-	require.Equal(t, d.GetOption(OptionBootfileName), []Option{bootFileOpt1, bootFileOpt2})
-	require.Equal(t, d.GetOption(OptionRouter), []Option{})
+	bootFileOpt := &OptBootfileName{"boot.img"}
+	d.UpdateOption(hostnameOpt)
+	d.UpdateOption(bootFileOpt)
 
 	require.Equal(t, d.GetOneOption(OptionHostName), hostnameOpt)
-	require.Equal(t, d.GetOneOption(OptionBootfileName), bootFileOpt1)
+	require.Equal(t, d.GetOneOption(OptionBootfileName), bootFileOpt)
 	require.Equal(t, d.GetOneOption(OptionRouter), nil)
 }
 
-func TestAddOption(t *testing.T) {
+func TestUpdateOption(t *testing.T) {
 	d, err := New()
 	require.NoError(t, err)
 
 	hostnameOpt := &OptionGeneric{OptionCode: OptionHostName, Data: []byte("darkstar")}
-	bootFileOpt1 := &OptionGeneric{OptionCode: OptionBootfileName, Data: []byte("boot.img")}
-	bootFileOpt2 := &OptionGeneric{OptionCode: OptionBootfileName, Data: []byte("boot2.img")}
-	d.AddOption(hostnameOpt)
-	d.AddOption(bootFileOpt1)
-	d.AddOption(bootFileOpt2)
+	bootFileOpt1 := &OptBootfileName{"boot.img"}
+	bootFileOpt2 := &OptBootfileName{"boot2.img"}
+	d.UpdateOption(hostnameOpt)
+	d.UpdateOption(bootFileOpt1)
+	d.UpdateOption(bootFileOpt2)
 
 	options := d.Options
-	require.Equal(t, len(options), 4)
-	require.Equal(t, options[3].Code(), OptionEnd)
-}
-
-func TestUpdateOption(t *testing.T) {
-	d, err := New()
-	require.NoError(t, err)
-	require.Equal(t, 1, len(d.Options))
-	require.Equal(t, OptionEnd, d.Options[0].Code())
-	// test that it will add the option since it's missing
-	d.UpdateOption(&OptDomainName{DomainName: "slackware.it"})
-	require.Equal(t, 2, len(d.Options))
-	require.Equal(t, OptionDomainName, d.Options[0].Code())
-	require.Equal(t, OptionEnd, d.Options[1].Code())
-	// test that it won't add another option of the same type
-	d.UpdateOption(&OptDomainName{DomainName: "slackware.it"})
-	require.Equal(t, 2, len(d.Options))
-	require.Equal(t, OptionDomainName, d.Options[0].Code())
-	require.Equal(t, OptionEnd, d.Options[1].Code())
+	require.Equal(t, len(options), 2)
+	require.Equal(t, d.Options.GetOne(OptionBootfileName), bootFileOpt2)
+	require.Equal(t, d.Options.GetOne(OptionHostName), hostnameOpt)
 }
 
 func TestDHCPv4NewRequestFromOffer(t *testing.T) {
 	offer, err := New()
 	require.NoError(t, err)
 	offer.SetBroadcast()
-	offer.AddOption(&OptMessageType{MessageType: MessageTypeOffer})
+	offer.UpdateOption(&OptMessageType{MessageType: MessageTypeOffer})
 	req, err := NewRequestFromOffer(offer)
 	require.Error(t, err)
 
 	// Now add the option so it doesn't error out.
-	offer.AddOption(&OptServerIdentifier{ServerID: net.IPv4(192, 168, 0, 1)})
+	offer.UpdateOption(&OptServerIdentifier{ServerID: net.IPv4(192, 168, 0, 1)})
 
 	// Broadcast request
 	req, err = NewRequestFromOffer(offer)
@@ -257,8 +235,8 @@ func TestDHCPv4NewRequestFromOffer(t *testing.T) {
 func TestDHCPv4NewRequestFromOfferWithModifier(t *testing.T) {
 	offer, err := New()
 	require.NoError(t, err)
-	offer.AddOption(&OptMessageType{MessageType: MessageTypeOffer})
-	offer.AddOption(&OptServerIdentifier{ServerID: net.IPv4(192, 168, 0, 1)})
+	offer.UpdateOption(&OptMessageType{MessageType: MessageTypeOffer})
+	offer.UpdateOption(&OptServerIdentifier{ServerID: net.IPv4(192, 168, 0, 1)})
 	userClass := WithUserClass([]byte("linuxboot"), false)
 	req, err := NewRequestFromOffer(offer, userClass)
 	require.NoError(t, err)
@@ -306,7 +284,6 @@ func TestNewDiscovery(t *testing.T) {
 	require.Equal(t, hwAddr, m.ClientHWAddr)
 	require.True(t, m.IsBroadcast())
 	require.True(t, m.Options.Has(OptionParameterRequestList))
-	require.True(t, m.Options.Has(OptionEnd))
 }
 
 func TestNewInform(t *testing.T) {
@@ -328,7 +305,7 @@ func TestIsOptionRequested(t *testing.T) {
 	require.False(t, pkt.IsOptionRequested(OptionDomainNameServer))
 
 	optprl := OptParameterRequestList{RequestedOpts: []OptionCode{OptionDomainNameServer}}
-	pkt.AddOption(&optprl)
+	pkt.UpdateOption(&optprl)
 	require.True(t, pkt.IsOptionRequested(OptionDomainNameServer))
 }