Add FLY_FORCE_INSTANCE support (#228)

Author: dangra

client.go

@@ -168,6 +168,7 @@ type ClientOptions struct {
 	Logger           Logger
 	EnableDebugTrace *bool
 	FlyForceRegion   *string
+	FlyForceInstance *string
 	Transport        *Transport
 }
 
@@ -206,6 +207,14 @@ func (t *Transport) setDefaults(opts *ClientOptions) {
 			t.FlyForceRegion = v
 		}
 	}
+
+	if opts.FlyForceInstance != nil {
+		t.FlyForceInstance = *opts.FlyForceInstance
+	} else if t.FlyForceInstance == "" {
+		if v := os.Getenv("FLY_FORCE_INSTANCE"); v != "" {
+			t.FlyForceInstance = v
+		}
+	}
 }
 
 func NewClientFromOptions(opts ClientOptions) *Client {
@@ -359,6 +368,7 @@ type Transport struct {
 	Tokens              *tokens.Tokens
 	EnableDebugTrace    bool
 	FlyForceRegion      string
+	FlyForceInstance    string
 }
 
 func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
@@ -372,6 +382,9 @@ func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
 	if t.FlyForceRegion != "" {
 		req.Header.Set("Fly-Force-Region", t.FlyForceRegion)
 	}
+	if t.FlyForceInstance != "" {
+		req.Header.Set("Fly-Force-Instance", t.FlyForceInstance)
+	}
 
 	return t.UnderlyingTransport.RoundTrip(req)
 }

client_test.go

@@ -83,6 +83,71 @@ func TestTransportSetDefaults_DoesNotOverrideFlyForceRegionFromTransport(t *test
 	}
 }
 
+func TestTransportSetDefaults_DoesNotOverrideFlyForceInstanceFromTransport(t *testing.T) {
+	t.Setenv("FLY_FORCE_INSTANCE", "worker-1")
+
+	transport := &Transport{FlyForceInstance: "worker-2"}
+	opts := ClientOptions{Transport: transport}
+
+	transport.setDefaults(&opts)
+
+	if transport.FlyForceInstance != "worker-2" {
+		t.Fatalf("expected FlyForceInstance to remain %q, got %q", "worker-2", transport.FlyForceInstance)
+	}
+}
+
+func TestTransportSetDefaults_SetsFlyForceInstanceFromEnv(t *testing.T) {
+	t.Setenv("FLY_FORCE_INSTANCE", "worker-1")
+
+	transport := &Transport{}
+	opts := ClientOptions{Transport: transport}
+
+	transport.setDefaults(&opts)
+
+	if transport.FlyForceInstance != "worker-1" {
+		t.Fatalf("expected FlyForceInstance to be %q, got %q", "worker-1", transport.FlyForceInstance)
+	}
+}
+
+type captureTripper struct {
+	req *http.Request
+}
+
+func (c *captureTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	c.req = req.Clone(req.Context())
+	c.req.Header = req.Header.Clone()
+	return &http.Response{
+		StatusCode: http.StatusOK,
+		Body:       io.NopCloser(strings.NewReader("{}")),
+		Header:     make(http.Header),
+	}, nil
+}
+
+func TestTransportRoundTrip_SetsFlyForceInstanceHeader(t *testing.T) {
+	capture := &captureTripper{}
+	transport := &Transport{
+		UnderlyingTransport: capture,
+		UserAgent:           "test/0",
+		Token:               "token",
+		FlyForceInstance:    "worker-2",
+	}
+
+	req, err := http.NewRequest(http.MethodGet, "http://example.test", nil)
+	if err != nil {
+		t.Fatalf("NewRequest returned error: %v", err)
+	}
+
+	resp, err := transport.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip returned error: %v", err)
+	}
+	defer resp.Body.Close()
+
+	if got := capture.req.Header.Get("Fly-Force-Instance"); got != "worker-2" {
+		t.Fatalf("Fly-Force-Instance header = %q, want %q", got, "worker-2")
+	}
+}
+
 func TestGraphQLOperationKind(t *testing.T) {
 	cases := []struct {
 		name string