1 files changed, 20 insertions, 28 deletions

diff --git a/dhcpv4/option_userclass_test.go b/dhcpv4/option_userclass_test.go
index e321a64..149fb92 100644
--- a/dhcpv4/option_userclass_test.go
+++ b/dhcpv4/option_userclass_test.go
@@ -7,11 +7,8 @@ import (
 )
 
 func TestOptUserClassToBytes(t *testing.T) {
-	opt := OptUserClass{
-		UserClasses: [][]byte{[]byte("linuxboot")},
-		Rfc3004:     true,
-	}
-	data := opt.ToBytes()
+	opt := OptRFC3004UserClass([][]byte{[]byte("linuxboot")})
+	data := opt.Value.ToBytes()
 	expected := []byte{
 		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
 	}
@@ -19,10 +16,8 @@ func TestOptUserClassToBytes(t *testing.T) {
 }
 
 func TestOptUserClassMicrosoftToBytes(t *testing.T) {
-	opt := OptUserClass{
-		UserClasses: [][]byte{[]byte("linuxboot")},
-	}
-	data := opt.ToBytes()
+	opt := OptUserClass([]byte("linuxboot"))
+	data := opt.Value.ToBytes()
 	expected := []byte{
 		'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
 	}
@@ -30,11 +25,12 @@ func TestOptUserClassMicrosoftToBytes(t *testing.T) {
 }
 
 func TestParseOptUserClassMultiple(t *testing.T) {
+	var opt UserClass
 	expected := []byte{
 		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
 		4, 't', 'e', 's', 't',
 	}
-	opt, err := ParseOptUserClass(expected)
+	err := opt.FromBytes(expected)
 	require.NoError(t, err)
 	require.Equal(t, len(opt.UserClasses), 2)
 	require.Equal(t, []byte("linuxboot"), opt.UserClasses[0])
@@ -42,50 +38,53 @@ func TestParseOptUserClassMultiple(t *testing.T) {
 }
 
 func TestParseOptUserClassNone(t *testing.T) {
+	var opt UserClass
 	expected := []byte{}
-	_, err := ParseOptUserClass(expected)
+	err := opt.FromBytes(expected)
 	require.Error(t, err)
 }
 
 func TestParseOptUserClassMicrosoft(t *testing.T) {
+	var opt UserClass
 	expected := []byte{
 		'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
 	}
-	opt, err := ParseOptUserClass(expected)
+	err := opt.FromBytes(expected)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(opt.UserClasses))
 	require.Equal(t, []byte("linuxboot"), opt.UserClasses[0])
 }
 
 func TestParseOptUserClassMicrosoftShort(t *testing.T) {
+	var opt UserClass
 	expected := []byte{
 		'l',
 	}
-	opt, err := ParseOptUserClass(expected)
+	err := opt.FromBytes(expected)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(opt.UserClasses))
 	require.Equal(t, []byte("l"), opt.UserClasses[0])
 }
 
 func TestParseOptUserClass(t *testing.T) {
+	var opt UserClass
 	expected := []byte{
 		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
 	}
-	opt, err := ParseOptUserClass(expected)
+	err := opt.FromBytes(expected)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(opt.UserClasses))
 	require.Equal(t, []byte("linuxboot"), opt.UserClasses[0])
 }
 
 func TestOptUserClassToBytesMultiple(t *testing.T) {
-	opt := OptUserClass{
-		UserClasses: [][]byte{
+	opt := OptRFC3004UserClass(
+		[][]byte{
 			[]byte("linuxboot"),
 			[]byte("test"),
 		},
-		Rfc3004: true,
-	}
-	data := opt.ToBytes()
+	)
+	data := opt.Value.ToBytes()
 	expected := []byte{
 		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
 		4, 't', 'e', 's', 't',
@@ -94,14 +93,7 @@ func TestOptUserClassToBytesMultiple(t *testing.T) {
 }
 
 func TestParseOptUserClassZeroLength(t *testing.T) {
-	expected := []byte{
-		0, 0,
-	}
-	_, err := ParseOptUserClass(expected)
+	var opt UserClass
+	err := opt.FromBytes([]byte{0, 0})
 	require.Error(t, err)
 }
-
-func TestOptUserClassCode(t *testing.T) {
-	opt := OptUserClass{}
-	require.Equal(t, OptionUserClassInformation, opt.Code())
-}