From c65f9a653e6f628cb4fd3f71e3c7a673e4d96fee Mon Sep 17 00:00:00 2001 From: mkorwel Date: Thu, 19 Mar 2026 00:04:34 -0700 Subject: [PATCH] feat(workspaces): complete task 1.2 with actual gce provisioning and unit tests --- .../src/services/computeService.test.ts | 61 +++++++++++++++ .../src/services/computeService.ts | 75 +++++++++++++++++-- .../src/services/workspaceService.test.ts | 65 ++++++++++++++++ plans/phase-1-workspace-core.md | 4 +- 4 files changed, 198 insertions(+), 7 deletions(-) create mode 100644 packages/workspace-manager/src/services/computeService.test.ts create mode 100644 packages/workspace-manager/src/services/workspaceService.test.ts diff --git a/packages/workspace-manager/src/services/computeService.test.ts b/packages/workspace-manager/src/services/computeService.test.ts new file mode 100644 index 0000000000..33bedaa16f --- /dev/null +++ b/packages/workspace-manager/src/services/computeService.test.ts @@ -0,0 +1,61 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { ComputeService } from './computeService.js'; + +const mockInstancesClient = { + insert: vi.fn().mockResolvedValue([{ latestResponse: { name: 'op-1' } }]), + delete: vi.fn().mockResolvedValue([{ latestResponse: { name: 'op-2' } }]), +}; + +vi.mock('@google-cloud/compute', () => ({ + InstancesClient: vi.fn().mockImplementation(() => mockInstancesClient), + })); + +describe('ComputeService', () => { + let service: ComputeService; + + beforeEach(() => { + vi.clearAllMocks(); + service = new ComputeService(); + }); + + describe('createWorkspaceInstance', () => { + it('should call instancesClient.insert with correct parameters', async () => { + const options = { + instanceName: 'test-inst', + machineType: 'e2-medium', + imageTag: 'latest', + zone: 'us-west1-a', + }; + + await service.createWorkspaceInstance(options); + + expect(mockInstancesClient.insert).toHaveBeenCalledWith( + expect.objectContaining({ + project: expect.any(String), + zone: 'us-west1-a', + instanceResource: expect.objectContaining({ + name: 'test-inst', + }), + }), + ); + }); + }); + + describe('deleteWorkspaceInstance', () => { + it('should call instancesClient.delete', async () => { + await service.deleteWorkspaceInstance('inst1', 'zone1'); + expect(mockInstancesClient.delete).toHaveBeenCalledWith( + expect.objectContaining({ + instance: 'inst1', + zone: 'zone1', + }), + ); + }); + }); +}); diff --git a/packages/workspace-manager/src/services/computeService.ts b/packages/workspace-manager/src/services/computeService.ts index 428e5cde44..f90bf6b087 100644 --- a/packages/workspace-manager/src/services/computeService.ts +++ b/packages/workspace-manager/src/services/computeService.ts @@ -15,20 +15,80 @@ export interface ProvisionOptions { export class ComputeService { private client: InstancesClient; + private projectId: string; constructor() { this.client = new InstancesClient(); + // In a real GCP environment, this is usually available via metadata server or env + this.projectId = process.env['GOOGLE_CLOUD_PROJECT'] || 'dev-project'; } /** * Provision a new GCE VM with the Workspace Container */ async createWorkspaceInstance(options: ProvisionOptions): Promise { - // TODO: Implement actual instancesClient.insert call - // For now, we just log and return + const { instanceName, machineType, imageTag, zone } = options; + + // TODO: Get the actual base image URL from config + const containerImage = `us-west1-docker.pkg.dev/${this.projectId}/workspaces/gemini-workspace:${imageTag}`; + + // The container declaration for Container-Optimized OS + const containerDeclaration = ` +spec: + containers: + - name: workspace + image: ${containerImage} + securityContext: + privileged: false + stdin: true + tty: true + restartPolicy: Always +`; + + const [operation] = await this.client.insert({ + project: this.projectId, + zone, + instanceResource: { + name: instanceName, + machineType: `zones/${zone}/machineTypes/${machineType}`, + disks: [ + { + boot: true, + autoDelete: true, + initializeParams: { + sourceImage: 'projects/cos-cloud/global/images/family/cos-stable', + }, + }, + ], + networkInterfaces: [ + { + network: 'global/networks/default', + // We use IAP for access, but a NAT might be needed for outbound internet + accessConfigs: [{ name: 'External NAT', type: 'ONE_TO_ONE_NAT' }], + }, + ], + metadata: { + items: [ + { + key: 'gce-container-declaration', + value: containerDeclaration, + }, + { + key: 'google-logging-enabled', + value: 'true', + }, + ], + }, + // Security: Tag for IAP access + tags: { + items: ['allow-ssh-iap'], + }, + }, + }); + // eslint-disable-next-line no-console console.log( - `[ComputeService] Mocking creation of ${options.instanceName} in ${options.zone}`, + `[ComputeService] Creation started for ${instanceName}. Op ID: ${operation.latestResponse.name}`, ); } @@ -39,10 +99,15 @@ export class ComputeService { instanceName: string, zone: string, ): Promise { - // TODO: Implement actual instancesClient.delete call + const [operation] = await this.client.delete({ + project: this.projectId, + zone, + instance: instanceName, + }); + // eslint-disable-next-line no-console console.log( - `[ComputeService] Mocking deletion of ${instanceName} in ${zone}`, + `[ComputeService] Deletion started for ${instanceName}. Op ID: ${operation.latestResponse.name}`, ); } } diff --git a/packages/workspace-manager/src/services/workspaceService.test.ts b/packages/workspace-manager/src/services/workspaceService.test.ts new file mode 100644 index 0000000000..8d63bb0fa4 --- /dev/null +++ b/packages/workspace-manager/src/services/workspaceService.test.ts @@ -0,0 +1,65 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { WorkspaceService } from './workspaceService.js'; +import { Firestore } from '@google-cloud/firestore'; + +const mockCollection = { + where: vi.fn().mockReturnThis(), + get: vi.fn(), + doc: vi.fn().mockReturnThis(), + set: vi.fn(), + delete: vi.fn(), +}; + +vi.mock('@google-cloud/firestore', () => ({ + Firestore: vi.fn().mockImplementation(() => ({ + collection: vi.fn().mockReturnValue(mockCollection), + })), + })); + +describe('WorkspaceService', () => { + let service: WorkspaceService; + + beforeEach(() => { + vi.clearAllMocks(); + service = new WorkspaceService(); + }); + + describe('listWorkspaces', () => { + it('should query Firestore for workspaces by ownerId', async () => { + const mockDocs = [{ id: '1', data: () => ({ name: 'ws1' }) }]; + mockCollection.get.mockResolvedValue({ docs: mockDocs }); + + const result = await service.listWorkspaces('user1'); + + expect(result).toHaveLength(1); + expect(result[0].id).toBe('1'); + expect(result[0].name).toBe('ws1'); + // Verify the collection name was 'workspaces' + const firestoreInstance = vi.mocked(Firestore).mock.results[0].value; + expect(firestoreInstance.collection).toHaveBeenCalledWith('workspaces'); + }); + }); + + describe('createWorkspace', () => { + it('should save workspace data to Firestore', async () => { + const data = { + owner_id: 'u1', + name: 'test', + instance_name: 'inst', + status: 'READY', + machine_type: 'e2', + zone: 'us1', + created_at: 'now', + }; + await service.createWorkspace('id1', data); + expect(mockCollection.doc).toHaveBeenCalledWith('id1'); + expect(mockCollection.set).toHaveBeenCalledWith(data); + }); + }); +}); diff --git a/plans/phase-1-workspace-core.md b/plans/phase-1-workspace-core.md index fc6953d1aa..fb19bf766b 100644 --- a/plans/phase-1-workspace-core.md +++ b/plans/phase-1-workspace-core.md @@ -25,8 +25,8 @@ Implement the core API to manage GCE-based workspaces. - [x] Initialize `packages/workspace-manager/`. - [x] Implement Express server for `/workspaces` (List, Create, Delete). - [x] Integrate Firestore to track workspace state (owner, instance_id, status). -- [ ] Integrate `@google-cloud/compute` for GCE instance lifecycle. -- [ ] Provision a VM with `Container-on-VM` settings pointing to the +- [x] Integrate `@google-cloud/compute` for GCE instance lifecycle. +- [x] Provision a VM with `Container-on-VM` settings pointing to the `gemini-workspace` image. ### Task 1.3: Cloud Run Deployment (v1)