github.com/GuanceCloud/cliutils@v1.1.21/pipeline/ptinput/ipdb/iploc/iploc_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 iploc
     7  
     8  import (
     9  	"io/fs"
    10  	"os"
    11  	"path/filepath"
    12  	"testing"
    13  
    14  	"github.com/ip2location/ip2location-go"
    15  	"github.com/stretchr/testify/assert"
    16  )
    17  
    18  var iplocationMockResult ip2location.IP2Locationrecord
    19  
    20  type mockDB struct{}
    21  
    22  func (db *mockDB) Get_all(ip string) (ip2location.IP2Locationrecord, error) { //nolint:stylecheck
    23  	return iplocationMockResult, nil
    24  }
    25  
    26  var mockOpenDB = func(f string) (DB, error) {
    27  	return &mockDB{}, nil
    28  }
    29  
    30  func TestIPloc(t *testing.T) {
    31  	openDB = mockOpenDB
    32  	iplocInstance := &IPloc{}
    33  	datadir, err := os.MkdirTemp(os.TempDir(), "iploc")
    34  	if err != nil {
    35  		t.Fatalf("tmp dir create error")
    36  	}
    37  	defer os.RemoveAll(datadir) //nolint:errcheck
    38  	iplocFile := "ip.bin"
    39  	ispFile := "isp.txt"
    40  
    41  	iplocDir := filepath.Join(datadir, "ipdb/iploc")
    42  	err = os.MkdirAll(iplocDir, os.ModePerm)
    43  	if err != nil {
    44  		l.Fatal(err)
    45  	}
    46  
    47  	err = os.WriteFile(filepath.Join(iplocDir, iplocFile), []byte(""), fs.ModePerm)
    48  	if err != nil {
    49  		t.Fatalf("iploc file create errror")
    50  	}
    51  
    52  	err = os.WriteFile(filepath.Join(iplocDir, ispFile), []byte("221.0.0.0/13 unicom"), fs.ModePerm)
    53  	if err != nil {
    54  		t.Fatalf("isp file create errror")
    55  	}
    56  
    57  	iplocInstance.Init(datadir, map[string]string{"isp_file": ispFile, "iploc_file": iplocFile})
    58  
    59  	t.Run("SearchIsp", func(t *testing.T) {
    60  		assert.Equal(t, "unicom", iplocInstance.SearchIsp("221.0.0.0"))
    61  		assert.Equal(t, "unknown", iplocInstance.SearchIsp("0.0.0.0"))
    62  	})
    63  
    64  	t.Run("Geo", func(t *testing.T) {
    65  		iplocationMockResult = ip2location.IP2Locationrecord{Country_short: "CN", City: ErrInvalidIP, Region: ErrInvalidIP}
    66  		ipInfo, err := iplocInstance.Geo("0.0.0.0")
    67  		assert.NoError(t, err)
    68  		assert.Equal(t, "CN", ipInfo.Country)
    69  		assert.Equal(t, "unknown", ipInfo.City)
    70  
    71  		backDB := iplocInstance.db
    72  		iplocInstance.db = nil
    73  		_, err = iplocInstance.Geo("0.0.0.0")
    74  
    75  		assert.NoError(t, err)
    76  
    77  		iplocInstance.db = backDB
    78  	})
    79  }