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 }