lib/types: Support MarshalText for NullIPPool

grafana · Apr 20, 2022 · b9879fb · b9879fb
1 parent da21ffb
commit b9879fb
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 9 deletions.
diff --git a/lib/options_test.go b/lib/options_test.go
@@ -473,17 +473,19 @@ func TestOptions(t *testing.T) {
 	})
 	t.Run("ClientIPRanges", func(t *testing.T) {
 		t.Parallel()
-		clientIPRanges, err := types.NewIPPool("129.112.232.12,123.12.0.0/32")
+		clientIPRanges := types.NullIPPool{}
+		err := clientIPRanges.UnmarshalText([]byte("129.112.232.12,123.12.0.0/32"))
 		require.NoError(t, err)
-		opts := Options{}.Apply(Options{LocalIPs: types.NullIPPool{Pool: clientIPRanges, Valid: true}})
+		opts := Options{}.Apply(Options{LocalIPs: clientIPRanges})
 		assert.NotNil(t, opts.LocalIPs)
 	})
 }
 
 func TestOptionsEnv(t *testing.T) {
 	t.Parallel()
-	mustIPPool := func(s string) *types.IPPool {
-		p, err := types.NewIPPool(s)
+	mustNullIPPool := func(s string) types.NullIPPool {
+		p := types.NullIPPool{}
+		err := p.UnmarshalText([]byte(s))
 		require.NoError(t, err)
 		return p
 	}
@@ -559,8 +561,8 @@ func TestOptionsEnv(t *testing.T) {
 		},
 		{"LocalIPs", "K6_LOCAL_IPS"}: {
 			"":                 types.NullIPPool{},
-			"192.168.220.2":    types.NullIPPool{Pool: mustIPPool("192.168.220.2"), Valid: true},
-			"192.168.220.2/24": types.NullIPPool{Pool: mustIPPool("192.168.220.0/24"), Valid: true},
+			"192.168.220.2":    mustNullIPPool("192.168.220.2"),
+			"192.168.220.0/24": mustNullIPPool("192.168.220.0/24"),
 		},
 		{"Throw", "K6_THROW"}: {
 			"":      null.Bool{},

diff --git a/lib/types/ipblock.go b/lib/types/ipblock.go
@@ -182,6 +182,7 @@ func (pool *IPPool) GetIPBig(index *big.Int) net.IP {
 type NullIPPool struct {
 	Pool  *IPPool
 	Valid bool
+	raw   []byte
 }
 
 // UnmarshalText converts text data to a valid NullIPPool
@@ -196,5 +197,11 @@ func (n *NullIPPool) UnmarshalText(data []byte) error {
 		return err
 	}
 	n.Valid = true
+	n.raw = data
 	return nil
 }
+
+// MarshalText returns the IPs pool in text form
+func (n *NullIPPool) MarshalText() ([]byte, error) {
+	return n.raw, nil
+}
diff --git a/lib/types/ipblock_test.go b/lib/types/ipblock_test.go
@@ -29,10 +29,8 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-//nolint:gochecknoglobals
-var max64 = new(big.Int).Exp(big.NewInt(2), big.NewInt(64), nil)
-
 func get128BigInt(hi, lo int64) *big.Int {
+	max64 := new(big.Int).Exp(big.NewInt(2), big.NewInt(64), nil)
 	h := big.NewInt(hi)
 	h.Mul(h, max64)
 	return h.Add(h, big.NewInt(lo))
@@ -160,3 +158,20 @@ func TestIpBlockError(t *testing.T) {
 		})
 	}
 }
+
+func TestNullIPPoolMarshalText(t *testing.T) {
+	t.Parallel()
+
+	rangeInput := "192.168.20.12-192.168.20.15,192.168.10.0/27"
+	p, err := NewIPPool(rangeInput)
+	require.NoError(t, err)
+
+	nullpool := NullIPPool{
+		Pool:  p,
+		Valid: true,
+		raw:   []byte(rangeInput),
+	}
+	text, err := nullpool.MarshalText()
+	require.NoError(t, err)
+	assert.Equal(t, "192.168.20.12-192.168.20.15,192.168.10.0/27", string(text))
+}