diff --git a/v1/providers/lambdalabs/common_test.go b/v1/providers/lambdalabs/common_test.go index 4c4259a4..a5d7b7de 100644 --- a/v1/providers/lambdalabs/common_test.go +++ b/v1/providers/lambdalabs/common_test.go @@ -9,6 +9,7 @@ import ( const ( testInstanceID = "test-instance-id" nonexistentInstance = "nonexistent-instance" + testPublicKey = "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ test@example.com" ) func setupMockClient() (*LambdaLabsClient, func()) { diff --git a/v1/providers/lambdalabs/gen/lambdalabs/model_ssh_key.go b/v1/providers/lambdalabs/gen/lambdalabs/model_ssh_key.go index d0504b8f..78260189 100644 --- a/v1/providers/lambdalabs/gen/lambdalabs/model_ssh_key.go +++ b/v1/providers/lambdalabs/gen/lambdalabs/model_ssh_key.go @@ -29,6 +29,8 @@ type SshKey struct { PublicKey string `json:"public_key"` // Private key for the SSH key. Only returned when generating a new key pair. PrivateKey NullableString `json:"private_key,omitempty"` + // Workspace ID associated with the SSH key. + WorkspaceId *string `json:"workspace_id,omitempty"` } type _SshKey SshKey @@ -184,6 +186,9 @@ func (o SshKey) ToMap() (map[string]interface{}, error) { if o.PrivateKey.IsSet() { toSerialize["private_key"] = o.PrivateKey.Get() } + if !IsNil(o.WorkspaceId) { + toSerialize["workspace_id"] = o.WorkspaceId + } return toSerialize, nil } diff --git a/v1/providers/lambdalabs/instance_test.go b/v1/providers/lambdalabs/instance_test.go index 6020659d..2f994002 100644 --- a/v1/providers/lambdalabs/instance_test.go +++ b/v1/providers/lambdalabs/instance_test.go @@ -2,6 +2,7 @@ package v1 import ( "context" + "encoding/json" "fmt" "testing" @@ -19,14 +20,16 @@ func TestLambdaLabsClient_CreateInstance_Success(t *testing.T) { defer cleanup() instanceID := testInstanceID - publicKey := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ test@example.com" + publicKey := testPublicKey + workspaceID := "ff393b16489f4f02838062b093181016" httpmock.RegisterResponder("POST", "https://cloud.lambda.ai/api/v1/ssh-keys", httpmock.NewJsonResponderOrPanic(200, openapi.AddSSHKey200Response{ Data: openapi.SshKey{ - Id: "ssh-key-id", - Name: "test-instance-id", - PublicKey: publicKey, + Id: "ssh-key-id", + Name: "test-instance-id", + PublicKey: publicKey, + WorkspaceId: &workspaceID, }, })) @@ -92,7 +95,7 @@ func TestLambdaLabsClient_CreateInstance_SSHKeyError(t *testing.T) { client, cleanup := setupMockClient() defer cleanup() - publicKey := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ test@example.com" + publicKey := testPublicKey httpmock.RegisterResponder("POST", "https://cloud.lambda.ai/api/v1/ssh-keys", httpmock.NewStringResponder(400, `{"error": {"code": "INVALID_REQUEST", "message": "SSH key already exists"}}`)) @@ -266,3 +269,39 @@ func TestLambdaLabsClient_RebootInstance_Error(t *testing.T) { err := client.RebootInstance(context.Background(), v1.CloudProviderInstanceID(instanceID)) assert.Error(t, err) } + +func TestLambdaLabsClient_CreateInstance_SSHKeyResponseWithWorkspaceID(t *testing.T) { + client, cleanup := setupMockClient() + defer cleanup() + + instanceID := testInstanceID + publicKey := testPublicKey + + // Raw JSON simulating Lambda Labs response after their Workspaces launch (June 2026). + // DisallowUnknownFields() in UnmarshalJSON would hard-error without WorkspaceId in the struct. + httpmock.RegisterResponder("POST", "https://cloud.lambda.ai/api/v1/ssh-keys", + httpmock.NewJsonResponderOrPanic(200, json.RawMessage(`{"data":{"id":"ssh-key-id","name":"test-instance-id","public_key":"ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ test@example.com","workspace_id":"ff393b16489f4f02838062b093181016"}}`))) + + httpmock.RegisterResponder("POST", "https://cloud.lambda.ai/api/v1/instance-operations/launch", + httpmock.NewJsonResponderOrPanic(200, openapi.LaunchInstance200Response{ + Data: openapi.LaunchInstance200ResponseData{ + InstanceIds: []string{instanceID}, + }, + })) + + mockInstance := createMockInstance(instanceID) + httpmock.RegisterResponder("GET", fmt.Sprintf("https://cloud.lambda.ai/api/v1/instances/%s", instanceID), + httpmock.NewJsonResponderOrPanic(200, openapi.GetInstance200Response{ + Data: mockInstance, + })) + + args := v1.CreateInstanceAttrs{ + InstanceType: "gpu_1x_a10", + Location: "us-west-1", + PublicKey: publicKey, + Name: "test-instance", + } + + _, err := client.CreateInstance(context.Background(), args) + require.NoError(t, err, "should not fail when Lambda Labs returns workspace_id in SSH key response") +}