test: collection: Add test for Assign

Signed-off-by: Daniel Xu <dxu@dxuuu.xyz>

collection_test.go

@@ -434,7 +434,7 @@ func TestCollectionSpec_LoadAndAssign_LazyLoading(t *testing.T) {
 	}
 }
 
-func TestCollectionAssign(t *testing.T) {
+func TestCollectionSpecAssign(t *testing.T) {
 	var specs struct {
 		Program *ProgramSpec `ebpf:"prog1"`
 		Map     *MapSpec     `ebpf:"map1"`
@@ -559,6 +559,89 @@ func TestAssignValues(t *testing.T) {
 
 }
 
+func TestCollectionAssign(t *testing.T) {
+	var objs struct {
+		Program *Program `ebpf:"prog1"`
+		Map     *Map     `ebpf:"map1"`
+	}
+
+	cs := &CollectionSpec{
+		Maps: map[string]*MapSpec{
+			"map1": {
+				Type:       Array,
+				KeySize:    4,
+				ValueSize:  4,
+				MaxEntries: 1,
+			},
+		},
+		Programs: map[string]*ProgramSpec{
+			"prog1": {
+				Type: SocketFilter,
+				Instructions: asm.Instructions{
+					asm.LoadImm(asm.R0, 0, asm.DWord),
+					asm.Return(),
+				},
+				License: "MIT",
+			},
+		},
+	}
+
+	coll, err := NewCollection(cs)
+	qt.Assert(t, err, qt.IsNil)
+	defer coll.Close()
+
+	qt.Assert(t, coll.Assign(&objs), qt.IsNil)
+	defer objs.Program.Close()
+	defer objs.Map.Close()
+
+	// Check that objs has received ownership of map and prog
+	qt.Assert(t, objs.Program.FD() >= 0, qt.IsTrue)
+	qt.Assert(t, objs.Map.FD() >= 0, qt.IsTrue)
+
+	// Check that the collection has lost ownership
+	qt.Assert(t, coll.Programs["prog1"], qt.IsNil)
+	qt.Assert(t, coll.Maps["map1"], qt.IsNil)
+}
+
+func TestCollectionAssignFail(t *testing.T) {
+	// `map2` does not exist
+	var objs struct {
+		Program *Program `ebpf:"prog1"`
+		Map     *Map     `ebpf:"map2"`
+	}
+
+	cs := &CollectionSpec{
+		Maps: map[string]*MapSpec{
+			"map1": {
+				Type:       Array,
+				KeySize:    4,
+				ValueSize:  4,
+				MaxEntries: 1,
+			},
+		},
+		Programs: map[string]*ProgramSpec{
+			"prog1": {
+				Type: SocketFilter,
+				Instructions: asm.Instructions{
+					asm.LoadImm(asm.R0, 0, asm.DWord),
+					asm.Return(),
+				},
+				License: "MIT",
+			},
+		},
+	}
+
+	coll, err := NewCollection(cs)
+	qt.Assert(t, err, qt.IsNil)
+	defer coll.Close()
+
+	qt.Assert(t, coll.Assign(&objs), qt.IsNotNil)
+
+	// Check that the collection has retained ownership
+	qt.Assert(t, coll.Programs["prog1"], qt.IsNotNil)
+	qt.Assert(t, coll.Maps["map1"], qt.IsNotNil)
+}
+
 func TestIncompleteLoadAndAssign(t *testing.T) {
 	spec := &CollectionSpec{
 		Programs: map[string]*ProgramSpec{