remove the getConnection method from PodmanConnection

containers · Feb 1, 2024 · fdbb930 · fdbb930
1 parent 286cef4
commit fdbb930
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 26 deletions.
diff --git a/packages/backend/src/managers/playground.spec.ts b/packages/backend/src/managers/playground.spec.ts
@@ -24,16 +24,19 @@ import type { PodmanConnection } from './podmanConnection';
 
 const mocks = vi.hoisted(() => ({
   postMessage: vi.fn(),
+  getContainerConnections: vi.fn(),
   pullImage: vi.fn(),
   createContainer: vi.fn(),
   stopContainer: vi.fn(),
   getFreePort: vi.fn(),
   containerRegistrySubscribeMock: vi.fn(),
-  getConnection: vi.fn(),
 }));
 
 vi.mock('@podman-desktop/api', async () => {
   return {
+    provider: {
+      getContainerConnections: mocks.getContainerConnections,
+    },
     containerEngine: {
       pullImage: mocks.pullImage,
       createContainer: mocks.createContainer,
@@ -62,27 +65,28 @@ beforeEach(() => {
       postMessage: mocks.postMessage,
     } as unknown as Webview,
     containerRegistryMock,
-    {
-      getConnection: mocks.getConnection,
-    } as unknown as PodmanConnection,
+    {} as PodmanConnection,
   );
 });
 
 test('startPlayground should fail if no provider', async () => {
   mocks.postMessage.mockResolvedValue(undefined);
+  mocks.getContainerConnections.mockReturnValue([]);
   await expect(manager.startPlayground('model1', '/path/to/model')).rejects.toThrowError(
     'Unable to find an engine to start playground',
   );
 });
 
 test('startPlayground should download image if not present then create container', async () => {
   mocks.postMessage.mockResolvedValue(undefined);
-  mocks.getConnection.mockReturnValue({
-    connection: {
-      type: 'podman',
-      status: () => 'started',
+  mocks.getContainerConnections.mockReturnValue([
+    {
+      connection: {
+        type: 'podman',
+        status: () => 'started',
+      },
     },
-  });
+  ]);
   vi.spyOn(manager, 'selectImage')
     .mockResolvedValueOnce(undefined)
     .mockResolvedValueOnce({
@@ -133,12 +137,14 @@ test('stopPlayground should fail if no playground is running', async () => {
 
 test('stopPlayground should stop a started playground', async () => {
   mocks.postMessage.mockResolvedValue(undefined);
-  mocks.getConnection.mockReturnValue({
-    connection: {
-      type: 'podman',
-      status: () => 'started',
+  mocks.getContainerConnections.mockReturnValue([
+    {
+      connection: {
+        type: 'podman',
+        status: () => 'started',
+      },
     },
-  });
+  ]);
   vi.spyOn(manager, 'selectImage').mockResolvedValue({
     Id: 'image1',
     engineId: 'engine1',

diff --git a/packages/backend/src/managers/playground.ts b/packages/backend/src/managers/playground.ts
@@ -16,7 +16,13 @@
  * SPDX-License-Identifier: Apache-2.0
  ***********************************************************************/
 
-import { containerEngine, type Webview, type ImageInfo } from '@podman-desktop/api';
+import {
+  containerEngine,
+  type Webview,
+  type ImageInfo,
+  type ProviderContainerConnection,
+  provider,
+} from '@podman-desktop/api';
 import type { LocalModelInfo } from '@shared/src/models/ILocalModelInfo';
 import type { ModelResponse } from '@shared/src/models/IModelResponse';
 
@@ -35,6 +41,14 @@ const LABEL_MODEL_PORT = 'ai-studio-model-port';
 // TODO: this should not be hardcoded
 const PLAYGROUND_IMAGE = 'quay.io/bootsy/playground:v0';
 
+function findFirstProvider(): ProviderContainerConnection | undefined {
+  const engines = provider
+    .getContainerConnections()
+    .filter(connection => connection.connection.type === 'podman')
+    .filter(connection => connection.connection.status() === 'started');
+  return engines.length > 0 ? engines[0] : undefined;
+}
+
 export class PlayGroundManager {
   private queryIdCounter = 0;
 
@@ -130,15 +144,15 @@ export class PlayGroundManager {
 
     this.setPlaygroundStatus(modelId, 'starting');
 
-    const connection = this.podmanConnection.getConnection();
+    const connection = findFirstProvider();
     if (!connection) {
       this.setPlaygroundStatus(modelId, 'error');
       throw new Error('Unable to find an engine to start playground');
     }
 
     let image = await this.selectImage(PLAYGROUND_IMAGE);
     if (!image) {
-      await containerEngine.pullImage(connection, PLAYGROUND_IMAGE, () => {});
+      await containerEngine.pullImage(connection.connection, PLAYGROUND_IMAGE, () => {});
       image = await this.selectImage(PLAYGROUND_IMAGE);
       if (!image) {
         this.setPlaygroundStatus(modelId, 'error');

diff --git a/packages/backend/src/managers/podmanConnection.ts b/packages/backend/src/managers/podmanConnection.ts
@@ -16,14 +16,12 @@
  * SPDX-License-Identifier: Apache-2.0
  ***********************************************************************/
 
-import { type ContainerProviderConnection, type RegisterContainerConnectionEvent, provider } from '@podman-desktop/api';
+import { type RegisterContainerConnectionEvent, provider } from '@podman-desktop/api';
 
 type startupHandle = () => void;
 
 export class PodmanConnection {
   #firstFound = false;
-  #connection: ContainerProviderConnection | undefined = undefined;
-
   #toExecuteAtStartup: startupHandle[] = [];
 
   init(): void {
@@ -37,7 +35,6 @@ export class PodmanConnection {
         return;
       }
       this.#firstFound = true;
-      this.#connection = e.connection;
       for (const f of this.#toExecuteAtStartup) {
         f();
       }
@@ -53,7 +50,6 @@ export class PodmanConnection {
     if (engines.length > 0) {
       disposable.dispose();
       this.#firstFound = true;
-      this.#connection = engines[0].connection;
     }
   }
 
@@ -66,8 +62,4 @@ export class PodmanConnection {
       this.#toExecuteAtStartup.push(f);
     }
   }
-
-  getConnection(): ContainerProviderConnection | undefined {
-    return this.#connection;
-  }
 }