|
| 1 | +package scaffold |
| 2 | + |
| 3 | +import ( |
| 4 | + "io" |
| 5 | + "net/http" |
| 6 | + "os" |
| 7 | + "path/filepath" |
| 8 | + "strings" |
| 9 | + |
| 10 | + "github.com/gruntwork-io/terratest/modules/k8s" |
| 11 | + . "github.com/onsi/gomega" |
| 12 | + "github.com/pkg/errors" |
| 13 | + corev1 "k8s.io/api/core/v1" |
| 14 | +) |
| 15 | + |
| 16 | +const ( |
| 17 | + _locustConfigMapTemplate = ` |
| 18 | +apiVersion: v1 |
| 19 | +kind: ConfigMap |
| 20 | +metadata: |
| 21 | + name: locust-config |
| 22 | +data: |
| 23 | + locustfile.py: |- |
| 24 | + from locust import HttpUser, task, between |
| 25 | +
|
| 26 | + class HttpbinRequester(HttpUser): |
| 27 | + @task |
| 28 | + def request_headers(self): |
| 29 | + self.client.get("/headers", headers={"Host": "httpbin.example"}) |
| 30 | + LOCUST_HOST: http://api7ee3-apisix-gateway-mtls:9080 |
| 31 | + LOCUST_SPAWN_RATE: "50" |
| 32 | + LOCUST_USERS: "500" |
| 33 | + LOCUST_AUTOSTART: "true" |
| 34 | +` |
| 35 | + _locustDeploymentTemplate = ` |
| 36 | +apiVersion: apps/v1 |
| 37 | +kind: Deployment |
| 38 | +metadata: |
| 39 | + name: locust |
| 40 | +spec: |
| 41 | + selector: |
| 42 | + matchLabels: |
| 43 | + app: locust |
| 44 | + template: |
| 45 | + metadata: |
| 46 | + labels: |
| 47 | + app: locust |
| 48 | + spec: |
| 49 | + containers: |
| 50 | + - name: locust |
| 51 | + image: locustio/locust |
| 52 | + ports: |
| 53 | + - containerPort: 8089 |
| 54 | + env: |
| 55 | + - name: LOCUST_HOST |
| 56 | + valueFrom: |
| 57 | + configMapKeyRef: |
| 58 | + name: locust-config |
| 59 | + key: LOCUST_HOST |
| 60 | + - name: LOCUST_SPAWN_RATE |
| 61 | + valueFrom: |
| 62 | + configMapKeyRef: |
| 63 | + name: locust-config |
| 64 | + key: LOCUST_SPAWN_RATE |
| 65 | + - name: LOCUST_USERS |
| 66 | + valueFrom: |
| 67 | + configMapKeyRef: |
| 68 | + name: locust-config |
| 69 | + key: LOCUST_USERS |
| 70 | + - name: LOCUST_AUTOSTART |
| 71 | + valueFrom: |
| 72 | + configMapKeyRef: |
| 73 | + name: locust-config |
| 74 | + key: LOCUST_AUTOSTART |
| 75 | + volumeMounts: |
| 76 | + - mountPath: /home/locust |
| 77 | + name: locust-config |
| 78 | + volumes: |
| 79 | + - name: locust-config |
| 80 | + configMap: |
| 81 | + name: locust-config |
| 82 | +` |
| 83 | + _locustServiceTemplate = ` |
| 84 | +apiVersion: v1 |
| 85 | +kind: Service |
| 86 | +metadata: |
| 87 | + name: locust |
| 88 | +spec: |
| 89 | + selector: |
| 90 | + app: locust |
| 91 | + ports: |
| 92 | + - name: web |
| 93 | + port: 8089 |
| 94 | + targetPort: 8089 |
| 95 | + protocol: TCP |
| 96 | + type: ClusterIP ` |
| 97 | +) |
| 98 | + |
| 99 | +func (s *Scaffold) DeployLocust() *corev1.Service { |
| 100 | + // create ConfigMap, Deployment, Service |
| 101 | + for _, yaml_ := range []string{_locustConfigMapTemplate, _locustDeploymentTemplate, _locustServiceTemplate} { |
| 102 | + err := s.CreateResourceFromString(yaml_) |
| 103 | + Expect(err).NotTo(HaveOccurred(), "create resource: %s", yaml_) |
| 104 | + } |
| 105 | + |
| 106 | + service, err := k8s.GetServiceE(s.t, s.kubectlOptions, "locust") |
| 107 | + Expect(err).NotTo(HaveOccurred(), "get service: locust") |
| 108 | + |
| 109 | + s.EnsureNumEndpointsReady(s.t, service.Name, 1) |
| 110 | + s.locustTunnel = k8s.NewTunnel(s.kubectlOptions, k8s.ResourceTypeService, "locust", 8089, 8089) |
| 111 | + s.addFinalizers(s.locustTunnel.Close) |
| 112 | + |
| 113 | + err = s.locustTunnel.ForwardPortE(s.t) |
| 114 | + Expect(err).NotTo(HaveOccurred(), "port-forward service: locust") |
| 115 | + |
| 116 | + return service |
| 117 | +} |
| 118 | + |
| 119 | +// func (s *Scaffold) LocustClient() *httpexpect.Expect { |
| 120 | +// u := url.URL{ |
| 121 | +// Scheme: "http", |
| 122 | +// Host: s.locustTunnel.Endpoint(), |
| 123 | +// } |
| 124 | +// return httpexpect.WithConfig(httpexpect.Config{ |
| 125 | +// BaseURL: u.String(), |
| 126 | +// Client: &http.Client{ |
| 127 | +// Transport: &http.Transport{}, |
| 128 | +// CheckRedirect: func(req *http.Request, via []*http.Request) error { |
| 129 | +// return http.ErrUseLastResponse |
| 130 | +// }, |
| 131 | +// }, |
| 132 | +// Reporter: httpexpect.NewAssertReporter( |
| 133 | +// httpexpect.NewAssertReporter(s.GinkgoT), |
| 134 | +// ), |
| 135 | +// }) |
| 136 | +// } |
| 137 | + |
| 138 | +func (s *Scaffold) ResetLocust() error { |
| 139 | + if s.locustTunnel == nil { |
| 140 | + return errors.New("locust is not deployed") |
| 141 | + } |
| 142 | + resp, err := http.Get("http://" + s.locustTunnel.Endpoint() + "/stats/reset") |
| 143 | + if err != nil { |
| 144 | + return errors.Wrap(err, "failed to request reset locust") |
| 145 | + } |
| 146 | + defer resp.Body.Close() |
| 147 | + if resp.StatusCode != http.StatusOK { |
| 148 | + return errors.Errorf("request reset locust not OK, status: %s", resp.Status) |
| 149 | + } |
| 150 | + return nil |
| 151 | +} |
| 152 | + |
| 153 | +func (s *Scaffold) DownloadLocustReport(filename string) error { |
| 154 | + if s.locustTunnel == nil { |
| 155 | + return errors.New("locust is not deployed") |
| 156 | + } |
| 157 | + if !strings.EqualFold(filepath.Ext(filename), ".html") { |
| 158 | + filename += ".html" |
| 159 | + } |
| 160 | + _ = os.Remove(filename) |
| 161 | + resp, err := http.Get("http://" + s.locustTunnel.Endpoint() + "/stats/report?download=1&theme=light") |
| 162 | + if err != nil { |
| 163 | + return errors.Wrap(err, "failed to request download report") |
| 164 | + } |
| 165 | + defer resp.Body.Close() |
| 166 | + if resp.StatusCode != http.StatusOK { |
| 167 | + return errors.Errorf("request download report not OK, status: %s", resp.Status) |
| 168 | + } |
| 169 | + data, err := io.ReadAll(resp.Body) |
| 170 | + if err != nil { |
| 171 | + return errors.Wrap(err, "failed to read report") |
| 172 | + } |
| 173 | + return os.WriteFile(filename, data, 0644) |
| 174 | +} |
0 commit comments