github.com/GuanceCloud/cliutils@v1.1.21/pipeline/ptinput/funcs/fn_gep_ip_test.go (about)

     1  // Unless explicitly stated otherwise all files in this repository are licensed
     2  // under the MIT License.
     3  // This product includes software developed at Guance Cloud (https://www.guance.com/).
     4  // Copyright 2021-present Guance, Inc.
     5  
     6  package funcs
     7  
     8  import (
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/GuanceCloud/cliutils/pipeline/ptinput"
    13  	"github.com/GuanceCloud/cliutils/pipeline/ptinput/ipdb"
    14  	"github.com/GuanceCloud/cliutils/point"
    15  
    16  	"github.com/stretchr/testify/assert"
    17  )
    18  
    19  type mockGEO struct{}
    20  
    21  func (m *mockGEO) Init(dataDir string, config map[string]string) {}
    22  func (m *mockGEO) SearchIsp(ip string) string                    { return geoDefaultVal }
    23  
    24  func (m *mockGEO) Geo(ip string) (*ipdb.IPdbRecord, error) {
    25  	return &ipdb.IPdbRecord{
    26  		City: func() string {
    27  			if ip == "unknown-city" {
    28  				return geoDefaultVal
    29  			} else {
    30  				return "Shanghai"
    31  			}
    32  		}(),
    33  		Region: func() string {
    34  			if ip == "unknown-region" {
    35  				return geoDefaultVal
    36  			} else {
    37  				return "Shanghai"
    38  			}
    39  		}(),
    40  		Country: func() string {
    41  			if ip == "unknown-country-short" {
    42  				return geoDefaultVal
    43  			} else {
    44  				return "CN"
    45  			}
    46  		}(),
    47  	}, nil
    48  }
    49  
    50  func TestGeoIpFunc(t *testing.T) {
    51  	cases := []struct {
    52  		in     string
    53  		script string
    54  
    55  		expected map[string]string
    56  
    57  		fail bool
    58  	}{
    59  		{
    60  			in: `{"ip":"1.2.3.4-something", "second":2,"third":"abc","forth":true}`,
    61  			script: `
    62  				json(_, ip)
    63  				geoip(ip)`,
    64  			expected: map[string]string{
    65  				"city":     "Shanghai",
    66  				"country":  "CN",
    67  				"province": "Shanghai",
    68  				"isp":      geoDefaultVal,
    69  			},
    70  		},
    71  
    72  		{
    73  			in: `{"ip":"unknown-city", "second":2,"third":"abc","forth":true}`,
    74  			script: `
    75  				json(_, ip)
    76  				geoip(ip)`,
    77  			expected: map[string]string{
    78  				"city":     geoDefaultVal,
    79  				"country":  "CN",
    80  				"province": "Shanghai",
    81  				"isp":      geoDefaultVal,
    82  			},
    83  		},
    84  
    85  		{
    86  			in: `{"aa": {"ip":"116.228.89.xxx"}, "second":2,"third":"abc","forth":true}`,
    87  			script: `
    88  				json(_, aa.ip)
    89  				geoip(aa.ip)`,
    90  			expected: map[string]string{
    91  				"city":     "Shanghai",
    92  				"country":  "CN",
    93  				"province": "Shanghai",
    94  				"isp":      geoDefaultVal,
    95  			},
    96  		},
    97  
    98  		{
    99  			in: `{"aa": {"ip":"unknown-region"}, "second":2,"third":"abc","forth":true}`,
   100  			script: `
   101  				json(_, aa.ip)
   102  				geoip(aa.ip)`,
   103  			expected: map[string]string{
   104  				"city":     "Shanghai",
   105  				"country":  "CN",
   106  				"province": geoDefaultVal,
   107  				"isp":      geoDefaultVal,
   108  			},
   109  		},
   110  
   111  		{
   112  			in: `{"aa": {"ip":"unknown-country-short"}, "second":2,"third":"abc","forth":true}`,
   113  			script: `
   114  				json(_, aa.ip)
   115  				geoip(aa.ip)`,
   116  			expected: map[string]string{
   117  				"city":     "Shanghai",
   118  				"country":  geoDefaultVal,
   119  				"province": "Shanghai",
   120  				"isp":      geoDefaultVal,
   121  			},
   122  		},
   123  	}
   124  
   125  	for idx, tc := range cases {
   126  		t.Logf("case %d...", idx)
   127  
   128  		runner, err := NewTestingRunner(tc.script)
   129  		if err != nil {
   130  			t.Errorf("[%d] failed: %s", idx, err)
   131  			return
   132  		}
   133  
   134  		pt := ptinput.NewPlPoint(
   135  			point.Logging, "test", nil, map[string]any{"message": tc.in}, time.Now())
   136  		pt.SetIPDB(&mockGEO{})
   137  		errR := runScript(runner, pt)
   138  
   139  		if errR != nil {
   140  			t.Fatal(errR.Error())
   141  		}
   142  
   143  		for k, v := range tc.expected {
   144  			r, _, e := pt.Get(k)
   145  			assert.NoError(t, e)
   146  			assert.Equal(t, v, r, "`%s` != `%s`, key: `%s`", r, v, k)
   147  		}
   148  	}
   149  }