github.com/kubewharf/katalyst-core@v0.5.3/pkg/custom-metric/collector/prometheus/collector_test.go (about)

     1  /*
     2  Copyright 2022 The Katalyst Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package prometheus
    18  
    19  import (
    20  	"bufio"
    21  	"context"
    22  	"net/http"
    23  	"net/http/httptest"
    24  	"os"
    25  	"path"
    26  	"strconv"
    27  	"strings"
    28  	"testing"
    29  
    30  	"github.com/stretchr/testify/assert"
    31  	v1 "k8s.io/api/core/v1"
    32  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    33  	"k8s.io/apimachinery/pkg/labels"
    34  
    35  	katalystbase "github.com/kubewharf/katalyst-core/cmd/base"
    36  	"github.com/kubewharf/katalyst-core/pkg/config/metric"
    37  	metricconf "github.com/kubewharf/katalyst-core/pkg/config/metric"
    38  	"github.com/kubewharf/katalyst-core/pkg/custom-metric/store/local"
    39  	"github.com/kubewharf/katalyst-core/pkg/util/general"
    40  	"github.com/kubewharf/katalyst-core/pkg/util/native"
    41  )
    42  
    43  var (
    44  	credentialPath = "/tmp/katalyst-ut/credential"
    45  	username       = "katalyst"
    46  	password       = "password"
    47  )
    48  
    49  func setupCredential(credentialPath string) error {
    50  	err := general.EnsureDirectory(credentialPath)
    51  	if err != nil {
    52  		return err
    53  	}
    54  
    55  	usernameFile, err := os.OpenFile(path.Join(credentialPath, fileNameUsername), os.O_WRONLY|os.O_CREATE, 0o755)
    56  	if err != nil {
    57  		return err
    58  	}
    59  	defer usernameFile.Close()
    60  	usernameWriter := bufio.NewWriter(usernameFile)
    61  	_, _ = usernameWriter.WriteString(username)
    62  	_ = usernameWriter.Flush()
    63  
    64  	passwordFile, err := os.OpenFile(path.Join(credentialPath, fileNamePassword), os.O_WRONLY|os.O_CREATE, 0o755)
    65  	if err != nil {
    66  		return err
    67  	}
    68  	defer passwordFile.Close()
    69  	passwordWriter := bufio.NewWriter(passwordFile)
    70  	_, _ = passwordWriter.WriteString(password)
    71  	_ = passwordWriter.Flush()
    72  
    73  	return nil
    74  }
    75  
    76  func TestPrometheusAddRequests(t *testing.T) {
    77  	t.Parallel()
    78  
    79  	ctx := context.Background()
    80  	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    81  		_, _ = w.Write([]byte(``))
    82  	}))
    83  	defer server.Close()
    84  
    85  	err2 := setupCredential(credentialPath)
    86  	assert.NoError(t, err2)
    87  
    88  	baseCtx, _ := katalystbase.GenerateFakeGenericContext(nil, nil, nil, nil)
    89  	genericConf := &metricconf.GenericMetricConfiguration{}
    90  	collectConf := &metric.CollectorConfiguration{
    91  		PodSelector:    labels.NewSelector(),
    92  		NodeSelector:   labels.NewSelector(),
    93  		CredentialPath: credentialPath,
    94  	}
    95  	storeConf := &metricconf.StoreConfiguration{}
    96  	localStore, _ := local.NewLocalMemoryMetricStore(ctx, baseCtx, genericConf, storeConf)
    97  
    98  	promCollector, err := NewPrometheusCollector(ctx, baseCtx, genericConf, collectConf, localStore)
    99  	assert.NoError(t, err)
   100  	promCollector.(*prometheusCollector).client, _ = newPrometheusClient()
   101  	assert.Equal(t, promCollector.(*prometheusCollector).username, username)
   102  	assert.Equal(t, promCollector.(*prometheusCollector).password, password)
   103  
   104  	hostAndPort := strings.Split(strings.TrimPrefix(server.URL, "http://"), ":")
   105  	assert.Equal(t, 2, len(hostAndPort))
   106  	port, _ := strconv.Atoi(hostAndPort[1])
   107  	promCollector.(*prometheusCollector).addRequest(&v1.Pod{
   108  		ObjectMeta: metav1.ObjectMeta{
   109  			Namespace: "ns1",
   110  			Name:      "name1",
   111  		},
   112  		Spec: v1.PodSpec{
   113  			Containers: []v1.Container{
   114  				{
   115  					Ports: []v1.ContainerPort{
   116  						{
   117  							Name:     native.ContainerMetricPortName,
   118  							HostPort: int32(port),
   119  						},
   120  					},
   121  				},
   122  			},
   123  		},
   124  		Status: v1.PodStatus{
   125  			HostIP: hostAndPort[0],
   126  		},
   127  	})
   128  	assert.Equal(t, 1, len(promCollector.(*prometheusCollector).scrapes))
   129  
   130  	promCollector.(*prometheusCollector).addRequest(&v1.Pod{
   131  		ObjectMeta: metav1.ObjectMeta{
   132  			Namespace: "ns1",
   133  			Name:      "name1",
   134  		},
   135  		Spec: v1.PodSpec{
   136  			Containers: []v1.Container{
   137  				{
   138  					Ports: []v1.ContainerPort{
   139  						{
   140  							Name:     native.ContainerMetricPortName,
   141  							HostPort: 11,
   142  						},
   143  					},
   144  				},
   145  			},
   146  		},
   147  		Status: v1.PodStatus{
   148  			HostIP: "not-exist",
   149  		},
   150  	})
   151  	assert.Equal(t, 1, len(promCollector.(*prometheusCollector).scrapes))
   152  }