1 files changed, 69 insertions, 0 deletions
diff --git a/dhcpv4/modifiers_test.go b/dhcpv4/modifiers_test.go
new file mode 100644
index 0000000..2f0e1ed
--- /dev/null
+++ b/dhcpv4/modifiers_test.go
@@ -0,0 +1,69 @@
+package dhcpv4
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestUserClassModifier(t *testing.T) {
+	d, _ := New()
+	userClass := WithUserClass([]byte("linuxboot"), false)
+	d = userClass(d)
+	expected := []byte{
+		77, // OptionUserClass
+		9,  // length
+		'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
+	}
+	require.Equal(t, "User Class Information -> linuxboot", d.options[0].String())
+	require.Equal(t, expected, d.options[0].ToBytes())
+}
+
+func TestUserClassModifierRFC(t *testing.T) {
+	d, _ := New()
+	userClass := WithUserClass([]byte("linuxboot"), true)
+	d = userClass(d)
+	expected := []byte{
+		77, // OptionUserClass
+		10, // length
+		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
+	}
+	require.Equal(t, "User Class Information -> linuxboot", d.options[0].String())
+	require.Equal(t, expected, d.options[0].ToBytes())
+}
+
+func TestWithNetboot(t *testing.T) {
+	d, _ := New()
+	d = WithNetboot(d)
+	require.Equal(t, "Parameter Request List -> [TFTP Server Name, Bootfile Name]", d.options[0].String())
+}
+
+func TestWithNetbootExistingTFTP(t *testing.T) {
+	d, _ := New()
+	OptParams := &OptParameterRequestList{
+		RequestedOpts: []OptionCode{OptionTFTPServerName},
+	}
+	d.AddOption(OptParams)
+	d = WithNetboot(d)
+	require.Equal(t, "Parameter Request List -> [TFTP Server Name, Bootfile Name]", d.options[0].String())
+}
+
+func TestWithNetbootExistingBootfileName(t *testing.T) {
+	d, _ := New()
+	OptParams := &OptParameterRequestList{
+		RequestedOpts: []OptionCode{OptionBootfileName},
+	}
+	d.AddOption(OptParams)
+	d = WithNetboot(d)
+	require.Equal(t, "Parameter Request List -> [Bootfile Name, TFTP Server Name]", d.options[0].String())
+}
+
+func TestWithNetbootExistingBoth(t *testing.T) {
+	d, _ := New()
+	OptParams := &OptParameterRequestList{
+		RequestedOpts: []OptionCode{OptionBootfileName, OptionTFTPServerName},
+	}
+	d.AddOption(OptParams)
+	d = WithNetboot(d)
+	require.Equal(t, "Parameter Request List -> [Bootfile Name, TFTP Server Name]", d.options[0].String())
+}