Skip to content

Commit 2f9dd19

Browse files
evacchievacchi
authored
refactor: migrate package-level flags to structured options (llm-d#283)
* refactor: migrate package-level flags to structured options Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com> * address comments Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com> --------- Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>
1 parent 23ce655 commit 2f9dd19

16 files changed

Lines changed: 1180 additions & 971 deletions

‎cmd/main.go‎

Lines changed: 11 additions & 415 deletions
Large diffs are not rendered by default.

‎cmd/main_test.go‎

Lines changed: 0 additions & 223 deletions
Original file line numberDiff line numberDiff line change
@@ -1,224 +1 @@
11
package main
2-
3-
import (
4-
"crypto/ecdsa"
5-
"crypto/elliptic"
6-
"crypto/rand"
7-
"crypto/tls"
8-
"crypto/x509"
9-
"crypto/x509/pkix"
10-
"encoding/pem"
11-
"math/big"
12-
"os"
13-
"path/filepath"
14-
"testing"
15-
"time"
16-
)
17-
18-
func generateTestCert(t *testing.T, dir string) (caPath, certPath, keyPath string) {
19-
t.Helper()
20-
21-
caKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
22-
if err != nil {
23-
t.Fatalf("generate CA key: %v", err)
24-
}
25-
caTemplate := &x509.Certificate{
26-
SerialNumber: big.NewInt(1),
27-
Subject: pkix.Name{CommonName: "Test CA"},
28-
NotBefore: time.Now().Add(-time.Hour),
29-
NotAfter: time.Now().Add(time.Hour),
30-
KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
31-
BasicConstraintsValid: true,
32-
IsCA: true,
33-
}
34-
caCertDER, err := x509.CreateCertificate(rand.Reader, caTemplate, caTemplate, &caKey.PublicKey, caKey)
35-
if err != nil {
36-
t.Fatalf("create CA cert: %v", err)
37-
}
38-
39-
caPath = filepath.Join(dir, "ca.pem")
40-
caPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: caCertDER})
41-
if err := os.WriteFile(caPath, caPEM, 0600); err != nil {
42-
t.Fatalf("write CA cert: %v", err)
43-
}
44-
45-
caCert, err := x509.ParseCertificate(caCertDER)
46-
if err != nil {
47-
t.Fatalf("parse CA cert: %v", err)
48-
}
49-
50-
clientKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
51-
if err != nil {
52-
t.Fatalf("generate client key: %v", err)
53-
}
54-
clientTemplate := &x509.Certificate{
55-
SerialNumber: big.NewInt(2),
56-
Subject: pkix.Name{CommonName: "Test Client"},
57-
NotBefore: time.Now().Add(-time.Hour),
58-
NotAfter: time.Now().Add(time.Hour),
59-
KeyUsage: x509.KeyUsageDigitalSignature,
60-
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
61-
}
62-
clientCertDER, err := x509.CreateCertificate(rand.Reader, clientTemplate, caCert, &clientKey.PublicKey, caKey)
63-
if err != nil {
64-
t.Fatalf("create client cert: %v", err)
65-
}
66-
67-
certPath = filepath.Join(dir, "client.pem")
68-
certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: clientCertDER})
69-
if err := os.WriteFile(certPath, certPEM, 0600); err != nil {
70-
t.Fatalf("write client cert: %v", err)
71-
}
72-
73-
keyPath = filepath.Join(dir, "client-key.pem")
74-
keyDER, err := x509.MarshalECPrivateKey(clientKey)
75-
if err != nil {
76-
t.Fatalf("marshal client key: %v", err)
77-
}
78-
keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})
79-
if err := os.WriteFile(keyPath, keyPEM, 0600); err != nil {
80-
t.Fatalf("write client key: %v", err)
81-
}
82-
83-
return caPath, certPath, keyPath
84-
}
85-
86-
func TestBuildTLSConfig_noFlags(t *testing.T) {
87-
cfg, err := buildTLSConfig("", "", "", false)
88-
if err != nil {
89-
t.Fatalf("unexpected error: %v", err)
90-
}
91-
if cfg != nil {
92-
t.Error("expected nil config when no TLS flags are set")
93-
}
94-
}
95-
96-
func TestBuildTLSConfig_mismatchedCertKey(t *testing.T) {
97-
tests := []struct {
98-
name string
99-
cert string
100-
key string
101-
}{
102-
{"cert without key", "/some/cert.pem", ""},
103-
{"key without cert", "", "/some/key.pem"},
104-
}
105-
for _, tt := range tests {
106-
t.Run(tt.name, func(t *testing.T) {
107-
_, err := buildTLSConfig("", tt.cert, tt.key, false)
108-
if err == nil {
109-
t.Fatal("expected error for mismatched cert/key")
110-
}
111-
})
112-
}
113-
}
114-
115-
func TestBuildTLSConfig_caCertOnly(t *testing.T) {
116-
dir := t.TempDir()
117-
caPath, _, _ := generateTestCert(t, dir)
118-
119-
cfg, err := buildTLSConfig(caPath, "", "", false)
120-
if err != nil {
121-
t.Fatalf("unexpected error: %v", err)
122-
}
123-
if cfg == nil {
124-
t.Fatal("expected non-nil TLS config")
125-
}
126-
if cfg.RootCAs == nil {
127-
t.Error("expected RootCAs to be populated")
128-
}
129-
if cfg.MinVersion != tls.VersionTLS12 {
130-
t.Errorf("MinVersion = %d, want %d", cfg.MinVersion, tls.VersionTLS12)
131-
}
132-
if len(cfg.Certificates) != 0 {
133-
t.Errorf("expected no client certificates, got %d", len(cfg.Certificates))
134-
}
135-
}
136-
137-
func TestBuildTLSConfig_mTLS(t *testing.T) {
138-
dir := t.TempDir()
139-
_, certPath, keyPath := generateTestCert(t, dir)
140-
141-
cfg, err := buildTLSConfig("", certPath, keyPath, false)
142-
if err != nil {
143-
t.Fatalf("unexpected error: %v", err)
144-
}
145-
if cfg == nil {
146-
t.Fatal("expected non-nil TLS config")
147-
}
148-
if len(cfg.Certificates) != 1 {
149-
t.Errorf("expected 1 client certificate, got %d", len(cfg.Certificates))
150-
}
151-
}
152-
153-
func TestBuildTLSConfig_caPlusMTLS(t *testing.T) {
154-
dir := t.TempDir()
155-
caPath, certPath, keyPath := generateTestCert(t, dir)
156-
157-
cfg, err := buildTLSConfig(caPath, certPath, keyPath, false)
158-
if err != nil {
159-
t.Fatalf("unexpected error: %v", err)
160-
}
161-
if cfg == nil {
162-
t.Fatal("expected non-nil TLS config")
163-
}
164-
if cfg.RootCAs == nil {
165-
t.Error("expected RootCAs to be populated")
166-
}
167-
if len(cfg.Certificates) != 1 {
168-
t.Errorf("expected 1 client certificate, got %d", len(cfg.Certificates))
169-
}
170-
}
171-
172-
func TestBuildTLSConfig_insecureSkipVerify(t *testing.T) {
173-
cfg, err := buildTLSConfig("", "", "", true)
174-
if err != nil {
175-
t.Fatalf("unexpected error: %v", err)
176-
}
177-
if cfg == nil {
178-
t.Fatal("expected non-nil TLS config")
179-
}
180-
if !cfg.InsecureSkipVerify {
181-
t.Error("expected InsecureSkipVerify to be true")
182-
}
183-
if cfg.MinVersion != tls.VersionTLS12 {
184-
t.Errorf("MinVersion = %d, want %d", cfg.MinVersion, tls.VersionTLS12)
185-
}
186-
}
187-
188-
func TestBuildTLSConfig_caCertFileNotFound(t *testing.T) {
189-
_, err := buildTLSConfig("/nonexistent/ca.pem", "", "", false)
190-
if err == nil {
191-
t.Fatal("expected error for missing CA cert file")
192-
}
193-
}
194-
195-
func TestBuildTLSConfig_malformedCACert(t *testing.T) {
196-
dir := t.TempDir()
197-
badCA := filepath.Join(dir, "bad-ca.pem")
198-
if err := os.WriteFile(badCA, []byte("not a valid PEM"), 0600); err != nil {
199-
t.Fatalf("write bad CA: %v", err)
200-
}
201-
202-
_, err := buildTLSConfig(badCA, "", "", false)
203-
if err == nil {
204-
t.Fatal("expected error for malformed CA cert")
205-
}
206-
}
207-
208-
func TestBuildTLSConfig_invalidCertKeyPair(t *testing.T) {
209-
dir := t.TempDir()
210-
211-
certPath := filepath.Join(dir, "cert.pem")
212-
keyPath := filepath.Join(dir, "key.pem")
213-
if err := os.WriteFile(certPath, []byte("not a cert"), 0600); err != nil {
214-
t.Fatalf("write bad cert: %v", err)
215-
}
216-
if err := os.WriteFile(keyPath, []byte("not a key"), 0600); err != nil {
217-
t.Fatalf("write bad key: %v", err)
218-
}
219-
220-
_, err := buildTLSConfig("", certPath, keyPath, false)
221-
if err == nil {
222-
t.Fatal("expected error for invalid cert/key pair")
223-
}
224-
}

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ require (
2222
github.com/prometheus/common v0.69.0
2323
github.com/redis/go-redis/extra/redisotel/v9 v9.21.0
2424
github.com/redis/go-redis/v9 v9.21.0
25+
github.com/spf13/pflag v1.0.10
2526
github.com/stretchr/testify v1.11.1
2627
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.69.0
2728
go.opentelemetry.io/otel v1.44.0
@@ -94,7 +95,6 @@ require (
9495
github.com/prometheus/procfs v0.17.0 // indirect
9596
github.com/redis/go-redis/extra/rediscmd/v9 v9.21.0 // indirect
9697
github.com/spf13/cobra v1.9.1 // indirect
97-
github.com/spf13/pflag v1.0.10 // indirect
9898
github.com/stoewer/go-strcase v1.3.0 // indirect
9999
github.com/x448/float16 v0.8.4 // indirect
100100
github.com/yuin/gopher-lua v1.1.1 // indirect

‎pkg/pubsub/options.go‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
package pubsub
2+
3+
import (
4+
"github.com/spf13/pflag"
5+
)
6+
7+
// Options holds CLI flags for the GCP PubSub flow.
8+
type Options struct {
9+
IGWBaseURL string
10+
ProjectID string
11+
RequestPathURL string
12+
InferenceObjective string
13+
RequestSubscriberID string
14+
ResultTopicID string
15+
TopicsConfigFile string
16+
BatchSize int
17+
}
18+
19+
func NewOptions() *Options {
20+
return &Options{
21+
RequestPathURL: "/v1/completions",
22+
BatchSize: 10,
23+
}
24+
}
25+
26+
func (o *Options) AddFlags(fs *pflag.FlagSet) {
27+
fs.StringVar(&o.IGWBaseURL, "pubsub.igw-base-url", o.IGWBaseURL, "Base URL for IGW. Mutually exclusive with pubsub.topics-config-file flag.")
28+
fs.StringVar(&o.ProjectID, "pubsub.project-id", o.ProjectID, "GCP project ID for PubSub")
29+
fs.StringVar(&o.RequestPathURL, "pubsub.request-path-url", o.RequestPathURL, "inference request path url. Mutually exclusive with pubsub.topics-config-file flag.")
30+
fs.StringVar(&o.InferenceObjective, "pubsub.inference-objective", o.InferenceObjective, "inference objective to use in requests. Mutually exclusive with pubsub.topics-config-file flag.")
31+
fs.StringVar(&o.RequestSubscriberID, "pubsub.request-subscriber-id", o.RequestSubscriberID, "GCP PubSub request topic subscriber ID. Mutually exclusive with pubsub.topics-config-file flag.")
32+
fs.StringVar(&o.ResultTopicID, "pubsub.result-topic-id", o.ResultTopicID, "GCP PubSub topic ID for results")
33+
fs.StringVar(&o.TopicsConfigFile, "pubsub.topics-config-file", o.TopicsConfigFile, "Topics Configuration file. Mutually exclusive with pubsub.igw-base-url, pubsub.request-subscriber-id, pubsub.request-path-url and pubsub.inference-objective flags. See documentation about syntax")
34+
fs.IntVar(&o.BatchSize, "pubsub.batch-size", o.BatchSize, "Number of inflight messages")
35+
}
36+
37+
// HasTopicConfig reports whether a multi-topic configuration file is set.
38+
func (o *Options) HasTopicConfig() bool {
39+
return o.TopicsConfigFile != ""
40+
}

0 commit comments

Comments
 (0)