update /contexts method

Author: mishushakov

js/tests/contexts.test.ts

@@ -5,9 +5,6 @@ import { sandboxTest } from './setup'
 sandboxTest('create context with no options', async ({ sandbox }) => {
   const context = await sandbox.createCodeContext()
 
-  // wait 1 second for the context to be created
-  await new Promise((resolve) => setTimeout(resolve, 1000))
-
   const contexts = await sandbox.listCodeContexts()
   const lastContext = contexts[contexts.length - 1]
 
@@ -22,9 +19,6 @@ sandboxTest('create context with options', async ({ sandbox }) => {
     cwd: '/root',
   })
 
-  // wait 1 second for the context to be created
-  await new Promise((resolve) => setTimeout(resolve, 1000))
-
   const contexts = await sandbox.listCodeContexts()
   const lastContext = contexts[contexts.length - 1]
 

python/tests/async/test_async_contexts.py

@@ -1,13 +1,9 @@
 from e2b_code_interpreter.code_interpreter_async import AsyncSandbox
-import asyncio
 
 
 async def test_create_context_with_no_options(async_sandbox: AsyncSandbox):
     context = await async_sandbox.create_code_context()
 
-    # wait 1 second for the context to be created
-    await asyncio.sleep(1)
-
     contexts = await async_sandbox.list_code_contexts()
     last_context = contexts[-1]
 
@@ -22,9 +18,6 @@ async def test_create_context_with_options(async_sandbox: AsyncSandbox):
         cwd="/root",
     )
 
-    # wait 1 second for the context to be created
-    await asyncio.sleep(1)
-
     contexts = await async_sandbox.list_code_contexts()
     last_context = contexts[-1]
 

python/tests/sync/test_contexts.py

@@ -1,13 +1,9 @@
 from e2b_code_interpreter.code_interpreter_sync import Sandbox
-import time
 
 
 def test_create_context_with_no_options(sandbox: Sandbox):
     context = sandbox.create_code_context()
 
-    # wait 1 second for the context to be created
-    time.sleep(1)
-
     contexts = sandbox.list_code_contexts()
     last_context = contexts[-1]
 
@@ -22,9 +18,6 @@ def test_create_context_with_options(sandbox: Sandbox):
         cwd="/root",
     )
 
-    # wait 1 second for the context to be created
-    time.sleep(1)
-
     contexts = sandbox.list_code_contexts()
     last_context = contexts[-1]
 

template/server/main.py

@@ -2,7 +2,7 @@
 import sys
 import httpx
 
-from typing import Dict, Union, Literal, Set
+from typing import Dict, Union, Literal, List
 
 from contextlib import asynccontextmanager
 from fastapi import FastAPI, Request
@@ -133,19 +133,18 @@ async def post_contexts(request: CreateContext) -> Context:
 
 
 @app.get("/contexts")
-async def get_contexts() -> Set[Context]:
+async def get_contexts() -> List[Context]:
     logger.info("Listing contexts")
 
-    context_ids = websockets.keys()
-
-    return set(
+    return [
         Context(
-            id=websockets[context_id].context_id,
-            language=websockets[context_id].language,
-            cwd=websockets[context_id].cwd,
+            id=ws.context_id,
+            language=ws.language,
+            cwd=ws.cwd,
         )
-        for context_id in context_ids
-    )
+        for key, ws in websockets.items()
+        if key != "default"
+    ]
 
 
 @app.post("/contexts/{context_id}/restart")