infra/packages/api/internal/orchestrator/nodemanager/node_test.go at 02cf6dd7f6f71d196d1beca45ecb5e035db279cc · e2b-dev/infra · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
package nodemanager_test

import (
	"context"
	"testing"

	"github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/e2b-dev/infra/packages/api/internal/api"
	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"
	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
)

func TestNode_OptimisticAdd_FlagEnabled(t *testing.T) {
	t.Parallel()

	// 1. Create a LaunchDarkly test data source
	td := ldtestdata.DataSource()

	// 2. Set the feature flag under test to true
	td.Update(td.Flag(featureflags.OptimisticResourceAccountingFlag.Key()).VariationForAll(true))

	// 3. Create a Feature Flag client with the test data source
	ffClient, err := featureflags.NewClientWithDatasource(td)
	require.NoError(t, err)

	// 4. Initialize Node with the injected ffClient
	node := nodemanager.NewTestNode("test-node", api.NodeStatusReady, 0, 4, nodemanager.WithFeatureFlags(ffClient))
	initialMetrics := node.Metrics()

	// 5. Call the method
	res := nodemanager.SandboxResources{
		CPUs:      2,
		MiBMemory: 1024,
	}
	node.OptimisticAdd(context.Background(), res)

	// 6. Assert: When flag is enabled, resources should be successfully accumulated
	newMetrics := node.Metrics()
	assert.Equal(t, initialMetrics.CpuAllocated+uint32(res.CPUs), newMetrics.CpuAllocated)
	assert.Equal(t, initialMetrics.MemoryAllocatedBytes+uint64(res.MiBMemory)*1024*1024, newMetrics.MemoryAllocatedBytes)
}

func TestNode_OptimisticAdd_FlagDisabled(t *testing.T) {
	t.Parallel()

	// 1. Create a LaunchDarkly test data source
	td := ldtestdata.DataSource()

	// 2. Set the feature flag under test to false
	td.Update(td.Flag(featureflags.OptimisticResourceAccountingFlag.Key()).VariationForAll(false))

	// 3. Create a Feature Flag client with the test data source
	ffClient, err := featureflags.NewClientWithDatasource(td)
	require.NoError(t, err)

	// 4. Initialize Node with the injected ffClient
	node := nodemanager.NewTestNode("test-node", api.NodeStatusReady, 0, 4, nodemanager.WithFeatureFlags(ffClient))
	initialMetrics := node.Metrics()

	// 5. Call the method
	res := nodemanager.SandboxResources{
		CPUs:      2,
		MiBMemory: 1024,
	}
	node.OptimisticAdd(context.Background(), res)

	// 6. Assert: When flag is disabled, return early, resources should not be accumulated
	newMetrics := node.Metrics()
	assert.Equal(t, initialMetrics.CpuAllocated, newMetrics.CpuAllocated)
	assert.Equal(t, initialMetrics.MemoryAllocatedBytes, newMetrics.MemoryAllocatedBytes)
}

func TestNode_OptimisticRemove_FlagEnabled(t *testing.T) {
	t.Parallel()

	// 1. Create a LaunchDarkly test data source
	td := ldtestdata.DataSource()

	// 2. Set the feature flag under test to true
	td.Update(td.Flag(featureflags.OptimisticResourceAccountingFlag.Key()).VariationForAll(true))

	// 3. Create a Feature Flag client with the test data source
	ffClient, err := featureflags.NewClientWithDatasource(td)
	require.NoError(t, err)

	// 4. Initialize Node with the injected ffClient - some resources are already allocated at initialization
	node := nodemanager.NewTestNode("test-node", api.NodeStatusReady, 4, 8192, nodemanager.WithFeatureFlags(ffClient))
	initialMetrics := node.Metrics()

	// 5. Call the method
	res := nodemanager.SandboxResources{
		CPUs:      2,
		MiBMemory: 1024,
	}
	node.OptimisticRemove(context.Background(), res)

	// 6. Assert: When flag is enabled, resources should be successfully deducted
	newMetrics := node.Metrics()
	assert.Equal(t, initialMetrics.CpuAllocated-uint32(res.CPUs), newMetrics.CpuAllocated)
	assert.Equal(t, initialMetrics.MemoryAllocatedBytes-uint64(res.MiBMemory)*1024*1024, newMetrics.MemoryAllocatedBytes)
}

func TestNode_OptimisticRemove_FlagDisabled(t *testing.T) {
	t.Parallel()

	// 1. Create a LaunchDarkly test data source
	td := ldtestdata.DataSource()

	// 2. Set the feature flag under test to false
	td.Update(td.Flag(featureflags.OptimisticResourceAccountingFlag.Key()).VariationForAll(false))

	// 3. Create a Feature Flag client with the test data source
	ffClient, err := featureflags.NewClientWithDatasource(td)
	require.NoError(t, err)

	// 4. Initialize Node with the injected ffClient - some resources are already allocated at initialization
	node := nodemanager.NewTestNode("test-node", api.NodeStatusReady, 4, 8192, nodemanager.WithFeatureFlags(ffClient))
	initialMetrics := node.Metrics()

	// 5. Call the method
	res := nodemanager.SandboxResources{
		CPUs:      2,
		MiBMemory: 1024,
	}
	node.OptimisticRemove(context.Background(), res)

	// 6. Assert: When flag is disabled, return early, resources should remain unchanged
	newMetrics := node.Metrics()
	assert.Equal(t, initialMetrics.CpuAllocated, newMetrics.CpuAllocated)
	assert.Equal(t, initialMetrics.MemoryAllocatedBytes, newMetrics.MemoryAllocatedBytes)
}