test: add missing tests for Entity+Edge model migration

models/record_test.go

@@ -34,3 +34,36 @@ func TestNewRecordWithEdges(t *testing.T) {
 	assert.Len(t, record.Edges(), 1)
 	assert.Equal(t, edge, record.Edges()[0])
 }
+
+func TestNewRecordWithMultipleEdges(t *testing.T) {
+	entity := &meteorv1beta1.Entity{
+		Urn:  "urn:test:scope:table:t1",
+		Name: "t1",
+		Type: "table",
+	}
+	lineageEdge := &meteorv1beta1.Edge{
+		SourceUrn: "urn:test:scope:table:t1",
+		TargetUrn: "urn:test:scope:table:t2",
+		Type:      "lineage",
+		Source:    "test",
+	}
+	ownerEdge := &meteorv1beta1.Edge{
+		SourceUrn: "urn:test:scope:table:t1",
+		TargetUrn: "urn:user:alice@co.com",
+		Type:      "owned_by",
+		Source:    "test",
+	}
+	record := models.NewRecord(entity, lineageEdge, ownerEdge)
+	assert.Equal(t, entity, record.Entity())
+	assert.Len(t, record.Edges(), 2)
+	assert.Equal(t, lineageEdge, record.Edges()[0])
+	assert.Equal(t, ownerEdge, record.Edges()[1])
+	assert.Equal(t, "lineage", record.Edges()[0].GetType())
+	assert.Equal(t, "owned_by", record.Edges()[1].GetType())
+}
+
+func TestNewRecordWithNilEntity(t *testing.T) {
+	record := models.NewRecord(nil)
+	assert.Nil(t, record.Entity())
+	assert.Empty(t, record.Edges())
+}

models/util_test.go

@@ -88,3 +88,66 @@ func TestRecordToJSON(t *testing.T) {
 	assert.Contains(t, string(b), `"entity"`)
 	assert.Contains(t, string(b), `"edges"`)
 }
+
+func TestNewEntityWithNilProps(t *testing.T) {
+	entity := models.NewEntity("urn:test:s:table:t1", "table", "t1", "test", nil)
+	assert.Equal(t, "urn:test:s:table:t1", entity.GetUrn())
+	assert.Equal(t, "table", entity.GetType())
+	assert.Nil(t, entity.GetProperties())
+}
+
+func TestNewEntityWithEmptyProps(t *testing.T) {
+	entity := models.NewEntity("urn:test:s:table:t1", "table", "t1", "test", map[string]any{})
+	assert.Equal(t, "urn:test:s:table:t1", entity.GetUrn())
+	assert.Nil(t, entity.GetProperties())
+}
+
+func TestNewEntityWithNestedProps(t *testing.T) {
+	entity := models.NewEntity("urn:test:s:table:t1", "table", "t1", "test", map[string]any{
+		"labels": map[string]string{"env": "production", "team": "data"},
+		"tags":   []string{"important", "verified"},
+	})
+	assert.Equal(t, "urn:test:s:table:t1", entity.GetUrn())
+	props := entity.GetProperties()
+	require.NotNil(t, props)
+	labelsVal := props.GetFields()["labels"].GetStructValue()
+	require.NotNil(t, labelsVal)
+	assert.Equal(t, "production", labelsVal.GetFields()["env"].GetStringValue())
+	assert.Equal(t, "data", labelsVal.GetFields()["team"].GetStringValue())
+	tagsVal := props.GetFields()["tags"].GetListValue()
+	require.NotNil(t, tagsVal)
+	assert.Len(t, tagsVal.GetValues(), 2)
+	assert.Equal(t, "important", tagsVal.GetValues()[0].GetStringValue())
+}
+
+func TestRecordToJSONWithoutEdges(t *testing.T) {
+	entity := &meteorv1beta1.Entity{
+		Urn:  "urn:test:s:table:t1",
+		Name: "t1",
+	}
+	record := models.NewRecord(entity)
+
+	b, err := models.RecordToJSON(record)
+	require.NoError(t, err)
+	assert.Contains(t, string(b), `"entity"`)
+	assert.NotContains(t, string(b), `"edges"`)
+}
+
+func TestRecordToJSONWithMultipleEdges(t *testing.T) {
+	entity := &meteorv1beta1.Entity{
+		Urn:  "urn:test:s:table:t1",
+		Name: "t1",
+	}
+	lineage := models.LineageEdge("urn:a", "urn:b", "test")
+	owner := models.OwnerEdge("urn:test:s:table:t1", "urn:user:bob@co.com", "test")
+	record := models.NewRecord(entity, lineage, owner)
+
+	b, err := models.RecordToJSON(record)
+	require.NoError(t, err)
+	s := string(b)
+	assert.Contains(t, s, `"entity"`)
+	assert.Contains(t, s, `"edges"`)
+	assert.Contains(t, s, `"lineage"`)
+	assert.Contains(t, s, `"owned_by"`)
+	assert.Contains(t, s, `"urn:user:bob@co.com"`)
+}

plugins/processors/enrich/processor_test.go

@@ -0,0 +1,135 @@
+//go:build plugins
+
+package enrich_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/raystack/meteor/models"
+	"github.com/raystack/meteor/plugins"
+	"github.com/raystack/meteor/plugins/processors/enrich"
+	testutils "github.com/raystack/meteor/test/utils"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestInit(t *testing.T) {
+	t.Run("should return error for invalid config", func(t *testing.T) {
+		proc := enrich.New(testutils.Logger)
+		err := proc.Init(context.Background(), plugins.Config{
+			RawConfig: map[string]any{},
+		})
+		assert.Error(t, err)
+	})
+
+	t.Run("should return no error for valid config", func(t *testing.T) {
+		proc := enrich.New(testutils.Logger)
+		err := proc.Init(context.Background(), plugins.Config{
+			RawConfig: map[string]any{
+				"attributes": map[string]any{
+					"team": "data-engineering",
+				},
+			},
+		})
+		assert.NoError(t, err)
+	})
+}
+
+func TestProcess(t *testing.T) {
+	t.Run("should enrich entity with nil properties", func(t *testing.T) {
+		proc := enrich.New(testutils.Logger)
+		err := proc.Init(context.Background(), plugins.Config{
+			RawConfig: map[string]any{
+				"attributes": map[string]any{
+					"team": "data-engineering",
+				},
+			},
+		})
+		require.NoError(t, err)
+
+		entity := models.NewEntity("urn:table:1", "table", "my-table", "bigquery", nil)
+		rec := models.NewRecord(entity)
+
+		result, err := proc.Process(context.Background(), rec)
+		require.NoError(t, err)
+
+		props := result.Entity().GetProperties().AsMap()
+		assert.Equal(t, "data-engineering", props["team"])
+	})
+
+	t.Run("should merge with existing properties", func(t *testing.T) {
+		proc := enrich.New(testutils.Logger)
+		err := proc.Init(context.Background(), plugins.Config{
+			RawConfig: map[string]any{
+				"attributes": map[string]any{
+					"team": "data-engineering",
+				},
+			},
+		})
+		require.NoError(t, err)
+
+		entity := models.NewEntity("urn:table:1", "table", "my-table", "bigquery", map[string]any{
+			"existing_key": "existing_value",
+		})
+		rec := models.NewRecord(entity)
+
+		result, err := proc.Process(context.Background(), rec)
+		require.NoError(t, err)
+
+		props := result.Entity().GetProperties().AsMap()
+		assert.Equal(t, "existing_value", props["existing_key"])
+		assert.Equal(t, "data-engineering", props["team"])
+	})
+
+	t.Run("should preserve edges through processing", func(t *testing.T) {
+		proc := enrich.New(testutils.Logger)
+		err := proc.Init(context.Background(), plugins.Config{
+			RawConfig: map[string]any{
+				"attributes": map[string]any{
+					"team": "data-engineering",
+				},
+			},
+		})
+		require.NoError(t, err)
+
+		entity := models.NewEntity("urn:table:1", "table", "my-table", "bigquery", nil)
+		lineage := models.LineageEdge("urn:table:1", "urn:table:2", "bigquery")
+		owner := models.OwnerEdge("urn:table:1", "urn:user:alice", "bigquery")
+		rec := models.NewRecord(entity, lineage, owner)
+
+		result, err := proc.Process(context.Background(), rec)
+		require.NoError(t, err)
+
+		assert.Len(t, result.Edges(), 2)
+		assert.Equal(t, lineage, result.Edges()[0])
+		assert.Equal(t, owner, result.Edges()[1])
+	})
+
+	t.Run("should only add string values and skip non-string values", func(t *testing.T) {
+		proc := enrich.New(testutils.Logger)
+		err := proc.Init(context.Background(), plugins.Config{
+			RawConfig: map[string]any{
+				"attributes": map[string]any{
+					"team":     "data-engineering",
+					"count":    42,
+					"enabled":  true,
+					"fraction": 3.14,
+				},
+			},
+		})
+		require.NoError(t, err)
+
+		entity := models.NewEntity("urn:table:1", "table", "my-table", "bigquery", nil)
+		rec := models.NewRecord(entity)
+
+		result, err := proc.Process(context.Background(), rec)
+		require.NoError(t, err)
+
+		props := result.Entity().GetProperties().AsMap()
+		assert.Equal(t, "data-engineering", props["team"])
+		assert.Nil(t, props["count"])
+		assert.Nil(t, props["enabled"])
+		assert.Nil(t, props["fraction"])
+	})
+}

plugins/processors/labels/labels_test.go

@@ -0,0 +1,147 @@
+//go:build plugins
+
+package labels_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/raystack/meteor/models"
+	meteorv1beta1 "github.com/raystack/meteor/models/raystack/meteor/v1beta1"
+	"github.com/raystack/meteor/plugins"
+	"github.com/raystack/meteor/plugins/processors/labels"
+	testutils "github.com/raystack/meteor/test/utils"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestInit(t *testing.T) {
+	t.Run("should return error for invalid config", func(t *testing.T) {
+		p := labels.New(testutils.Logger)
+		err := p.Init(context.Background(), plugins.Config{
+			RawConfig: map[string]any{},
+		})
+		assert.Error(t, err)
+	})
+
+	t.Run("should return no error for valid config", func(t *testing.T) {
+		p := labels.New(testutils.Logger)
+		err := p.Init(context.Background(), plugins.Config{
+			RawConfig: map[string]any{
+				"labels": map[string]any{
+					"team": "data-eng",
+				},
+			},
+		})
+		assert.NoError(t, err)
+	})
+}
+
+func TestProcess(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("should add labels to entity with nil properties", func(t *testing.T) {
+		p := labels.New(testutils.Logger)
+		require.NoError(t, p.Init(ctx, plugins.Config{
+			RawConfig: map[string]any{
+				"labels": map[string]any{
+					"team": "data-eng",
+				},
+			},
+		}))
+
+		entity := models.NewEntity("urn:test:scope:table:t1", "table", "t1", "test", nil)
+		rec := models.NewRecord(entity)
+
+		result, err := p.Process(ctx, rec)
+		require.NoError(t, err)
+
+		props := result.Entity().GetProperties().AsMap()
+		lbls, ok := props["labels"].(map[string]any)
+		require.True(t, ok)
+		assert.Equal(t, "data-eng", lbls["team"])
+	})
+
+	t.Run("should add labels to entity with existing properties but no labels", func(t *testing.T) {
+		p := labels.New(testutils.Logger)
+		require.NoError(t, p.Init(ctx, plugins.Config{
+			RawConfig: map[string]any{
+				"labels": map[string]any{
+					"env": "production",
+				},
+			},
+		}))
+
+		entity := models.NewEntity("urn:test:scope:table:t2", "table", "t2", "test", map[string]any{
+			"description": "some table",
+		})
+		rec := models.NewRecord(entity)
+
+		result, err := p.Process(ctx, rec)
+		require.NoError(t, err)
+
+		props := result.Entity().GetProperties().AsMap()
+		assert.Equal(t, "some table", props["description"])
+		lbls, ok := props["labels"].(map[string]any)
+		require.True(t, ok)
+		assert.Equal(t, "production", lbls["env"])
+	})
+
+	t.Run("should merge labels with existing labels", func(t *testing.T) {
+		p := labels.New(testutils.Logger)
+		require.NoError(t, p.Init(ctx, plugins.Config{
+			RawConfig: map[string]any{
+				"labels": map[string]any{
+					"env":  "production",
+					"team": "data-eng",
+				},
+			},
+		}))
+
+		entity := models.NewEntity("urn:test:scope:table:t3", "table", "t3", "test", map[string]any{
+			"labels": map[string]any{
+				"existing": "value",
+				"team":     "old-team",
+			},
+		})
+		rec := models.NewRecord(entity)
+
+		result, err := p.Process(ctx, rec)
+		require.NoError(t, err)
+
+		props := result.Entity().GetProperties().AsMap()
+		lbls, ok := props["labels"].(map[string]any)
+		require.True(t, ok)
+		assert.Equal(t, "value", lbls["existing"])
+		assert.Equal(t, "data-eng", lbls["team"])
+		assert.Equal(t, "production", lbls["env"])
+	})
+
+	t.Run("should preserve edges through processing", func(t *testing.T) {
+		p := labels.New(testutils.Logger)
+		require.NoError(t, p.Init(ctx, plugins.Config{
+			RawConfig: map[string]any{
+				"labels": map[string]any{
+					"team": "data-eng",
+				},
+			},
+		}))
+
+		entity := models.NewEntity("urn:test:scope:table:t4", "table", "t4", "test", nil)
+		edges := []*meteorv1beta1.Edge{
+			models.OwnerEdge("urn:test:scope:table:t4", "urn:test:scope:user:alice", "test"),
+			models.LineageEdge("urn:test:scope:table:t4", "urn:test:scope:table:t5", "test"),
+		}
+		rec := models.NewRecord(entity, edges...)
+
+		result, err := p.Process(ctx, rec)
+		require.NoError(t, err)
+
+		resultEdges := result.Edges()
+		require.Len(t, resultEdges, 2)
+		assert.Equal(t, "owned_by", resultEdges[0].Type)
+		assert.Equal(t, "urn:test:scope:user:alice", resultEdges[0].TargetUrn)
+		assert.Equal(t, "lineage", resultEdges[1].Type)
+		assert.Equal(t, "urn:test:scope:table:t5", resultEdges[1].TargetUrn)
+	})
+}