11package nodemanager_test
22
33import (
4+ "context"
45 "testing"
56
7+ "github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata"
68 "github.com/stretchr/testify/assert"
9+ "github.com/stretchr/testify/require"
710
811 "github.com/e2b-dev/infra/packages/api/internal/api"
912 "github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"
13+ "github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
1014)
1115
12- func TestNode_OptimisticAdd (t * testing.T ) {
16+ func TestNode_OptimisticAdd_FlagEnabled (t * testing.T ) {
1317 t .Parallel ()
14-
15- // Initialize an idle node
16- node := nodemanager .NewTestNode ("test-node" , api .NodeStatusReady , 0 , 4 )
18+
19+ // 1. Create a LaunchDarkly test data source
20+ td := ldtestdata .DataSource ()
21+
22+ // 2. Set the feature flag under test to true
23+ td .Update (td .Flag (featureflags .OptimisticResourceAccountingFlag .Key ()).VariationForAll (true ))
24+
25+ // 3. Create a Feature Flag client with the test data source
26+ ffClient , err := featureflags .NewClientWithDatasource (td )
27+ require .NoError (t , err )
28+
29+ // 4. Initialize Node with the injected ffClient
30+ node := nodemanager .NewTestNode ("test-node" , api .NodeStatusReady , 0 , 4 , nodemanager .WithFeatureFlags (ffClient ))
1731 initialMetrics := node .Metrics ()
1832
19- // Simulate the resources to be allocated
33+ // 5. Call the method
2034 res := nodemanager.SandboxResources {
2135 CPUs : 2 ,
22- MiBMemory : 1024 , // 1GB
36+ MiBMemory : 1024 ,
2337 }
38+ node .OptimisticAdd (context .Background (), res )
2439
25- // Perform optimistic addition
26- node .OptimisticAdd (res )
40+ // 6. Assert: When flag is enabled, resources should be successfully accumulated
2741 newMetrics := node .Metrics ()
28-
29- // Verify that CPU and memory are increased as expected
3042 assert .Equal (t , initialMetrics .CpuAllocated + uint32 (res .CPUs ), newMetrics .CpuAllocated )
3143 assert .Equal (t , initialMetrics .MemoryAllocatedBytes + uint64 (res .MiBMemory )* 1024 * 1024 , newMetrics .MemoryAllocatedBytes )
32- }
44+ }
45+
46+ func TestNode_OptimisticAdd_FlagDisabled (t * testing.T ) {
47+ t .Parallel ()
48+
49+ // 1. Create a LaunchDarkly test data source
50+ td := ldtestdata .DataSource ()
51+
52+ // 2. Set the feature flag under test to false
53+ td .Update (td .Flag (featureflags .OptimisticResourceAccountingFlag .Key ()).VariationForAll (false ))
54+
55+ // 3. Create a Feature Flag client with the test data source
56+ ffClient , err := featureflags .NewClientWithDatasource (td )
57+ require .NoError (t , err )
58+
59+ // 4. Initialize Node with the injected ffClient
60+ node := nodemanager .NewTestNode ("test-node" , api .NodeStatusReady , 0 , 4 , nodemanager .WithFeatureFlags (ffClient ))
61+ initialMetrics := node .Metrics ()
62+
63+ // 5. Call the method
64+ res := nodemanager.SandboxResources {
65+ CPUs : 2 ,
66+ MiBMemory : 1024 ,
67+ }
68+ node .OptimisticAdd (context .Background (), res )
69+
70+ // 6. Assert: When flag is disabled, return early, resources should not be accumulated
71+ newMetrics := node .Metrics ()
72+ assert .Equal (t , initialMetrics .CpuAllocated , newMetrics .CpuAllocated )
73+ assert .Equal (t , initialMetrics .MemoryAllocatedBytes , newMetrics .MemoryAllocatedBytes )
74+ }
75+
76+ func TestNode_OptimisticRemove_FlagEnabled (t * testing.T ) {
77+ t .Parallel ()
78+
79+ // 1. Create a LaunchDarkly test data source
80+ td := ldtestdata .DataSource ()
81+
82+ // 2. Set the feature flag under test to true
83+ td .Update (td .Flag (featureflags .OptimisticResourceAccountingFlag .Key ()).VariationForAll (true ))
84+
85+ // 3. Create a Feature Flag client with the test data source
86+ ffClient , err := featureflags .NewClientWithDatasource (td )
87+ require .NoError (t , err )
88+
89+ // 4. Initialize Node with the injected ffClient - some resources are already allocated at initialization
90+ node := nodemanager .NewTestNode ("test-node" , api .NodeStatusReady , 4 , 8192 , nodemanager .WithFeatureFlags (ffClient ))
91+ initialMetrics := node .Metrics ()
92+
93+ // 5. Call the method
94+ res := nodemanager.SandboxResources {
95+ CPUs : 2 ,
96+ MiBMemory : 1024 ,
97+ }
98+ node .OptimisticRemove (context .Background (), res )
99+
100+ // 6. Assert: When flag is enabled, resources should be successfully deducted
101+ newMetrics := node .Metrics ()
102+ assert .Equal (t , initialMetrics .CpuAllocated - uint32 (res .CPUs ), newMetrics .CpuAllocated )
103+ assert .Equal (t , initialMetrics .MemoryAllocatedBytes - uint64 (res .MiBMemory )* 1024 * 1024 , newMetrics .MemoryAllocatedBytes )
104+ }
105+
106+ func TestNode_OptimisticRemove_FlagDisabled (t * testing.T ) {
107+ t .Parallel ()
108+
109+ // 1. Create a LaunchDarkly test data source
110+ td := ldtestdata .DataSource ()
111+
112+ // 2. Set the feature flag under test to false
113+ td .Update (td .Flag (featureflags .OptimisticResourceAccountingFlag .Key ()).VariationForAll (false ))
114+
115+ // 3. Create a Feature Flag client with the test data source
116+ ffClient , err := featureflags .NewClientWithDatasource (td )
117+ require .NoError (t , err )
118+
119+ // 4. Initialize Node with the injected ffClient - some resources are already allocated at initialization
120+ node := nodemanager .NewTestNode ("test-node" , api .NodeStatusReady , 4 , 8192 , nodemanager .WithFeatureFlags (ffClient ))
121+ initialMetrics := node .Metrics ()
122+
123+ // 5. Call the method
124+ res := nodemanager.SandboxResources {
125+ CPUs : 2 ,
126+ MiBMemory : 1024 ,
127+ }
128+ node .OptimisticRemove (context .Background (), res )
129+
130+ // 6. Assert: When flag is disabled, return early, resources should remain unchanged
131+ newMetrics := node .Metrics ()
132+ assert .Equal (t , initialMetrics .CpuAllocated , newMetrics .CpuAllocated )
133+ assert .Equal (t , initialMetrics .MemoryAllocatedBytes , newMetrics .MemoryAllocatedBytes )
134+ }
0 commit comments