unit tests

Signed-off-by: MorningTZH <morningtzh@yeah.net>

Author: morningtzh

packages/api/internal/orchestrator/nodemanager/node_test.go

@@ -0,0 +1,32 @@
+package nodemanager_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"
+)
+
+func TestNode_OptimisticAdd(t *testing.T) {
+	t.Parallel()
+	
+	// Initialize an idle node
+	node := nodemanager.NewTestNode("test-node", api.NodeStatusReady, 0, 4)
+	initialMetrics := node.Metrics()
+
+	// Simulate the resources to be allocated
+	res := nodemanager.SandboxResources{
+		CPUs:      2,
+		MiBMemory: 1024, // 1GB
+	}
+
+	// Perform optimistic addition
+	node.OptimisticAdd(res)
+	newMetrics := node.Metrics()
+
+	// Verify that CPU and memory are increased as expected
+	assert.Equal(t, initialMetrics.CpuAllocated+uint32(res.CPUs), newMetrics.CpuAllocated)
+	assert.Equal(t, initialMetrics.MemoryAllocatedBytes+uint64(res.MiBMemory)*1024*1024, newMetrics.MemoryAllocatedBytes)
+}

packages/api/internal/orchestrator/placement/placement_best_of_K_test.go

@@ -65,6 +65,36 @@ func TestBestOfK_Score_PreferBiggerNode(t *testing.T) {
 	assert.Greater(t, score, score2)
 }
 
+func TestBestOfK_Score_WithPendingResources(t *testing.T) {
+	t.Parallel()
+	config := DefaultBestOfKConfig()
+	algo := NewBestOfK(config).(*BestOfK)
+
+	// Create two nodes with identical base loads
+	nodeNormal := nodemanager.NewTestNode("node-normal", api.NodeStatusReady, 0, 4)
+	nodeWithPending := nodemanager.NewTestNode("node-pending", api.NodeStatusReady, 0, 4)
+
+	// Inject InProgress resources into nodeWithPending using StartPlacing
+	// This simulates a Sandbox that is currently being placed but hasn't fully started
+	pendingRes := nodemanager.SandboxResources{
+		CPUs:      2,
+		MiBMemory: 1024,
+	}
+	nodeWithPending.PlacementMetrics.StartPlacing("pending-sbx-1", pendingRes)
+
+	reqResources := nodemanager.SandboxResources{
+		CPUs:      1,
+		MiBMemory: 512,
+	}
+
+	scoreNormal := algo.Score(nodeNormal, reqResources, config)
+	scorePending := algo.Score(nodeWithPending, reqResources, config)
+
+	// A node with pending resources has a higher 'reserved' CPU count, 
+	// so its calculated Score should be greater (meaning worse/lower priority)
+	assert.Greater(t, scorePending, scoreNormal, "Node with pending resources should receive a higher (worse) score")
+}
+
 func TestBestOfK_CanFit(t *testing.T) {
 	t.Parallel()
 	config := DefaultBestOfKConfig()

packages/api/internal/orchestrator/placement/placement_test.go

@@ -213,3 +213,37 @@ func TestPlaceSandbox_ResourceExhausted(t *testing.T) {
 	// Verify node1 was NOT excluded (ResourceExhausted nodes should be retried)
 	algorithm.AssertNumberOfCalls(t, "chooseNode", 2)
 }
+
+func TestPlaceSandbox_TriggersOptimisticUpdate(t *testing.T) {
+	t.Parallel()
+	ctx := t.Context()
+
+	// Create a node and record the initial allocated CPU
+	node1 := nodemanager.NewTestNode("node1", api.NodeStatusReady, 0, 4)
+	initialCpuAllocated := node1.Metrics().CpuAllocated
+
+	nodes := []*nodemanager.Node{node1}
+	
+	// Mock algorithm directly returns node1
+	algorithm := &mockAlgorithm{}
+	algorithm.On("chooseNode", mock.Anything, nodes, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+		Return(node1, nil)
+
+	// Request 2 vCPUs
+	sbxRequest := &orchestrator.SandboxCreateRequest{
+		Sandbox: &orchestrator.SandboxConfig{
+			SandboxId: "test-optimistic-sandbox",
+			Vcpu:      2,
+			RamMb:     1024,
+		},
+	}
+
+	resultNode, err := PlaceSandbox(ctx, algorithm, nodes, nil, sbxRequest, machineinfo.MachineInfo{}, false, nil)
+
+	require.NoError(t, err)
+	assert.NotNil(t, resultNode)
+	
+	// Verify: After successful placement, the node's CpuAllocated should be increased by 2 from the base
+	updatedCpuAllocated := resultNode.Metrics().CpuAllocated
+	assert.Equal(t, initialCpuAllocated+2, updatedCpuAllocated, "Node metrics should be optimistically updated after placement")
+}