More Roundtrip tests

2025-12-17 22:52:24 +02:00
parent 5e1448dcdb
commit a427aa5537
23 changed files with 22897 additions and 1319 deletions
--- a/pkg/writers/dbml/writer.go
+++ b/pkg/writers/dbml/writer.go
@@ -29,7 +29,6 @@ func (w *Writer) WriteDatabase(db *models.Database) error {
 		return os.WriteFile(w.options.OutputPath, []byte(content), 0644)
 	}

-	// If no output path, print to stdout
 	fmt.Print(content)
 	return nil
 }
@@ -48,7 +47,7 @@ func (w *Writer) WriteSchema(schema *models.Schema) error {

 // WriteTable writes a Table model to DBML format
 func (w *Writer) WriteTable(table *models.Table) error {
-	content := w.tableToDBML(table, table.Schema)
+	content := w.tableToDBML(table)

 	if w.options.OutputPath != "" {
 		return os.WriteFile(w.options.OutputPath, []byte(content), 0644)
@@ -60,70 +59,63 @@ func (w *Writer) WriteTable(table *models.Table) error {

 // databaseToDBML converts a Database to DBML format string
 func (w *Writer) databaseToDBML(d *models.Database) string {
-	var result string
+	var sb strings.Builder

-	// Add database comment if exists
 	if d.Description != "" {
-		result += fmt.Sprintf("// %s\n", d.Description)
+		sb.WriteString(fmt.Sprintf("// %s\n", d.Description))
 	}
 	if d.Comment != "" {
-		result += fmt.Sprintf("// %s\n", d.Comment)
+		sb.WriteString(fmt.Sprintf("// %s\n", d.Comment))
 	}
 	if d.Description != "" || d.Comment != "" {
-		result += "\n"
+		sb.WriteString("\n")
 	}

-	// Process each schema
 	for _, schema := range d.Schemas {
-		result += w.schemaToDBML(schema)
+		sb.WriteString(w.schemaToDBML(schema))
 	}

-	// Add relationships
-	result += "\n// Relationships\n"
+	sb.WriteString("\n// Relationships\n")
 	for _, schema := range d.Schemas {
 		for _, table := range schema.Tables {
 			for _, constraint := range table.Constraints {
 				if constraint.Type == models.ForeignKeyConstraint {
-					result += w.constraintToDBML(constraint, schema.Name, table.Name)
+					sb.WriteString(w.constraintToDBML(constraint, table))
 				}
 			}
 		}
 	}

-	return result
+	return sb.String()
 }

 // schemaToDBML converts a Schema to DBML format string
 func (w *Writer) schemaToDBML(schema *models.Schema) string {
-	var result string
+	var sb strings.Builder

 	if schema.Description != "" {
-		result += fmt.Sprintf("// Schema: %s - %s\n", schema.Name, schema.Description)
+		sb.WriteString(fmt.Sprintf("// Schema: %s - %s\n", schema.Name, schema.Description))
 	}

-	// Process tables
 	for _, table := range schema.Tables {
-		result += w.tableToDBML(table, schema.Name)
-		result += "\n"
+		sb.WriteString(w.tableToDBML(table))
+		sb.WriteString("\n")
 	}

-	return result
+	return sb.String()
 }

 // tableToDBML converts a Table to DBML format string
-func (w *Writer) tableToDBML(t *models.Table, schemaName string) string {
-	var result string
+func (w *Writer) tableToDBML(t *models.Table) string {
+	var sb strings.Builder

-	// Table definition
-	tableName := fmt.Sprintf("%s.%s", schemaName, t.Name)
-	result += fmt.Sprintf("Table %s {\n", tableName)
+	tableName := fmt.Sprintf("%s.%s", t.Schema, t.Name)
+	sb.WriteString(fmt.Sprintf("Table %s {\n", tableName))

-	// Add columns
 	for _, column := range t.Columns {
-		result += fmt.Sprintf("  %s %s", column.Name, column.Type)
+		sb.WriteString(fmt.Sprintf("  %s %s", column.Name, column.Type))

-		// Add column attributes
-		attrs := make([]string, 0)
+		var attrs []string
 		if column.IsPrimaryKey {
 			attrs = append(attrs, "pk")
 		}
@@ -134,77 +126,74 @@ func (w *Writer) tableToDBML(t *models.Table, schemaName string) string {
 			attrs = append(attrs, "increment")
 		}
 		if column.Default != nil {
-			attrs = append(attrs, fmt.Sprintf("default: %v", column.Default))
+			attrs = append(attrs, fmt.Sprintf("default: `%v`", column.Default))
 		}

 		if len(attrs) > 0 {
-			result += fmt.Sprintf(" [%s]", strings.Join(attrs, ", "))
+			sb.WriteString(fmt.Sprintf(" [%s]", strings.Join(attrs, ", ")))
 		}

 		if column.Comment != "" {
-			result += fmt.Sprintf(" // %s", column.Comment)
+			sb.WriteString(fmt.Sprintf(" // %s", column.Comment))
 		}
-		result += "\n"
+		sb.WriteString("\n")
 	}

-	// Add indexes
-	indexCount := 0
-	for _, index := range t.Indexes {
-		if indexCount == 0 {
-			result += "\n  indexes {\n"
-		}
-		indexAttrs := make([]string, 0)
-		if index.Unique {
-			indexAttrs = append(indexAttrs, "unique")
-		}
-		if index.Name != "" {
-			indexAttrs = append(indexAttrs, fmt.Sprintf("name: '%s'", index.Name))
-		}
-		if index.Type != "" {
-			indexAttrs = append(indexAttrs, fmt.Sprintf("type: %s", index.Type))
-		}
+	if len(t.Indexes) > 0 {
+		sb.WriteString("\n  indexes {\n")
+		for _, index := range t.Indexes {
+			var indexAttrs []string
+			if index.Unique {
+				indexAttrs = append(indexAttrs, "unique")
+			}
+			if index.Name != "" {
+				indexAttrs = append(indexAttrs, fmt.Sprintf("name: '%s'", index.Name))
+			}
+			if index.Type != "" {
+				indexAttrs = append(indexAttrs, fmt.Sprintf("type: %s", index.Type))
+			}

-		result += fmt.Sprintf("    (%s)", strings.Join(index.Columns, ", "))
-		if len(indexAttrs) > 0 {
-			result += fmt.Sprintf(" [%s]", strings.Join(indexAttrs, ", "))
+			sb.WriteString(fmt.Sprintf("    (%s)", strings.Join(index.Columns, ", ")))
+			if len(indexAttrs) > 0 {
+				sb.WriteString(fmt.Sprintf(" [%s]", strings.Join(indexAttrs, ", ")))
+			}
+			sb.WriteString("\n")
 		}
-		result += "\n"
-		indexCount++
-	}
-	if indexCount > 0 {
-		result += "  }\n"
+		sb.WriteString("  }\n")
 	}

-	// Add table note
-	if t.Description != "" || t.Comment != "" {
-		note := t.Description
-		if note != "" && t.Comment != "" {
-			note += " - "
-		}
-		note += t.Comment
-		result += fmt.Sprintf("\n  Note: '%s'\n", note)
+	note := strings.TrimSpace(t.Description + " " + t.Comment)
+	if note != "" {
+		sb.WriteString(fmt.Sprintf("\n  Note: '%s'\n", note))
 	}

-	result += "}\n"
-	return result
+	sb.WriteString("}\n")
+	return sb.String()
 }

 // constraintToDBML converts a Constraint to DBML format string
-func (w *Writer) constraintToDBML(c *models.Constraint, schemaName, tableName string) string {
+func (w *Writer) constraintToDBML(c *models.Constraint, t *models.Table) string {
 	if c.Type != models.ForeignKeyConstraint || c.ReferencedTable == "" {
 		return ""
 	}

-	fromTable := fmt.Sprintf("%s.%s", schemaName, tableName)
+	fromTable := fmt.Sprintf("%s.%s", c.Schema, c.Table)
 	toTable := fmt.Sprintf("%s.%s", c.ReferencedSchema, c.ReferencedTable)

-	// Determine relationship cardinality
-	// For foreign keys, it's typically many-to-one
-	relationship := ">"
+	relationship := ">" // Default to many-to-one
+	for _, index := range t.Indexes {
+		if index.Unique && strings.Join(index.Columns, ",") == strings.Join(c.Columns, ",") {
+			relationship = "-" // one-to-one
+			break
+		}
+	}
+	for _, column := range c.Columns {
+		if t.Columns[column].IsPrimaryKey {
+			relationship = "-" // one-to-one
+			break
+		}
+	}

-	// Build from and to column references
-	// For single columns: table.column
-	// For multiple columns: table.(col1, col2)
 	var fromRef, toRef string
 	if len(c.Columns) == 1 {
 		fromRef = fmt.Sprintf("%s.%s", fromTable, c.Columns[0])
@@ -218,20 +207,18 @@ func (w *Writer) constraintToDBML(c *models.Constraint, schemaName, tableName st
 		toRef = fmt.Sprintf("%s.(%s)", toTable, strings.Join(c.ReferencedColumns, ", "))
 	}

-	result := fmt.Sprintf("Ref: %s %s %s", fromRef, relationship, toRef)
-
-	// Add actions
-	actions := make([]string, 0)
+	var actions []string
 	if c.OnDelete != "" {
-		actions = append(actions, fmt.Sprintf("ondelete: %s", c.OnDelete))
+		actions = append(actions, fmt.Sprintf("delete: %s", c.OnDelete))
 	}
 	if c.OnUpdate != "" {
-		actions = append(actions, fmt.Sprintf("onupdate: %s", c.OnUpdate))
-	}
-	if len(actions) > 0 {
-		result += fmt.Sprintf(" [%s]", strings.Join(actions, ", "))
+		actions = append(actions, fmt.Sprintf("update: %s", c.OnUpdate))
 	}

-	result += "\n"
-	return result
-}
+	refLine := fmt.Sprintf("Ref: %s %s %s", fromRef, relationship, toRef)
+	if len(actions) > 0 {
+		refLine += fmt.Sprintf(" [%s]", strings.Join(actions, ", "))
+	}
+
+	return refLine + "\n"
+}
--- a/pkg/writers/dbml/writer_test.go
+++ b/pkg/writers/dbml/writer_test.go
@@ -3,11 +3,11 @@ package dbml
 import (
 	"os"
 	"path/filepath"
-	"strings"
 	"testing"

 	"git.warky.dev/wdevs/relspecgo/pkg/models"
 	"git.warky.dev/wdevs/relspecgo/pkg/writers"
+	"github.com/stretchr/testify/assert"
 )

 func TestWriter_WriteTable(t *testing.T) {
@@ -46,96 +46,40 @@ func TestWriter_WriteTable(t *testing.T) {

 	writer := NewWriter(opts)
 	err := writer.WriteTable(table)
-	if err != nil {
-		t.Fatalf("WriteTable() error = %v", err)
-	}
+	assert.NoError(t, err)

 	content, err := os.ReadFile(outputPath)
-	if err != nil {
-		t.Fatalf("Failed to read output file: %v", err)
-	}
-
+	assert.NoError(t, err)
 	output := string(content)

-	// Verify table structure
-	if !strings.Contains(output, "Table public.users {") {
-		t.Error("Output should contain table definition")
-	}
-
-	// Verify columns
-	if !strings.Contains(output, "id bigint") {
-		t.Error("Output should contain id column")
-	}
-	if !strings.Contains(output, "pk") {
-		t.Error("Output should contain pk attribute for id")
-	}
-	if !strings.Contains(output, "increment") {
-		t.Error("Output should contain increment attribute for id")
-	}
-	if !strings.Contains(output, "email varchar(255)") {
-		t.Error("Output should contain email column")
-	}
-	if !strings.Contains(output, "not null") {
-		t.Error("Output should contain not null attribute")
-	}
-
-	// Verify table note
-	if !strings.Contains(output, "Note:") && table.Description != "" {
-		t.Error("Output should contain table note when description is present")
-	}
+	assert.Contains(t, output, "Table public.users {")
+	assert.Contains(t, output, "id bigint [pk, increment]")
+	assert.Contains(t, output, "email varchar(255) [not null]")
+	assert.Contains(t, output, "Note: 'User accounts table'")
 }

 func TestWriter_WriteDatabase_WithRelationships(t *testing.T) {
 	db := models.InitDatabase("test_db")
 	schema := models.InitSchema("public")

-	// Create users table
 	usersTable := models.InitTable("users", "public")
 	idCol := models.InitColumn("id", "users", "public")
 	idCol.Type = "bigint"
 	idCol.IsPrimaryKey = true
-	idCol.AutoIncrement = true
-	idCol.NotNull = true
 	usersTable.Columns["id"] = idCol
-
-	emailCol := models.InitColumn("email", "users", "public")
-	emailCol.Type = "varchar(255)"
-	emailCol.NotNull = true
-	usersTable.Columns["email"] = emailCol
-
-	// Add index to users table
-	emailIdx := models.InitIndex("idx_users_email")
+	emailIdx := models.InitIndex("idx_users_email", "users", "public")
 	emailIdx.Columns = []string{"email"}
 	emailIdx.Unique = true
-	emailIdx.Table = "users"
-	emailIdx.Schema = "public"
 	usersTable.Indexes["idx_users_email"] = emailIdx
+	schema.Tables = append(schema.Tables, usersTable)

-	// Create posts table
 	postsTable := models.InitTable("posts", "public")
 	postIdCol := models.InitColumn("id", "posts", "public")
 	postIdCol.Type = "bigint"
-	postIdCol.IsPrimaryKey = true
-	postIdCol.AutoIncrement = true
-	postIdCol.NotNull = true
 	postsTable.Columns["id"] = postIdCol
-
 	userIdCol := models.InitColumn("user_id", "posts", "public")
 	userIdCol.Type = "bigint"
-	userIdCol.NotNull = true
 	postsTable.Columns["user_id"] = userIdCol
-
-	titleCol := models.InitColumn("title", "posts", "public")
-	titleCol.Type = "varchar(200)"
-	titleCol.NotNull = true
-	postsTable.Columns["title"] = titleCol
-
-	publishedCol := models.InitColumn("published", "posts", "public")
-	publishedCol.Type = "boolean"
-	publishedCol.Default = "false"
-	postsTable.Columns["published"] = publishedCol
-
-	// Add foreign key constraint
 	fk := models.InitConstraint("fk_posts_user", models.ForeignKeyConstraint)
 	fk.Table = "posts"
 	fk.Schema = "public"
@@ -144,353 +88,68 @@ func TestWriter_WriteDatabase_WithRelationships(t *testing.T) {
 	fk.ReferencedSchema = "public"
 	fk.ReferencedColumns = []string{"id"}
 	fk.OnDelete = "CASCADE"
-	fk.OnUpdate = "CASCADE"
 	postsTable.Constraints["fk_posts_user"] = fk
-
-	schema.Tables = append(schema.Tables, usersTable, postsTable)
+	schema.Tables = append(schema.Tables, postsTable)
 	db.Schemas = append(db.Schemas, schema)

 	tmpDir := t.TempDir()
 	outputPath := filepath.Join(tmpDir, "test.dbml")
-
-	opts := &writers.WriterOptions{
-		OutputPath: outputPath,
-	}
-
+	opts := &writers.WriterOptions{OutputPath: outputPath}
 	writer := NewWriter(opts)
 	err := writer.WriteDatabase(db)
-	if err != nil {
-		t.Fatalf("WriteDatabase() error = %v", err)
-	}
+	assert.NoError(t, err)

 	content, err := os.ReadFile(outputPath)
-	if err != nil {
-		t.Fatalf("Failed to read output file: %v", err)
-	}
-
+	assert.NoError(t, err)
 	output := string(content)

-	// Verify tables
-	if !strings.Contains(output, "Table public.users {") {
-		t.Error("Output should contain users table")
-	}
-	if !strings.Contains(output, "Table public.posts {") {
-		t.Error("Output should contain posts table")
-	}
-
-	// Verify foreign key reference
-	if !strings.Contains(output, "Ref:") {
-		t.Error("Output should contain Ref for foreign key")
-	}
-	if !strings.Contains(output, "public.posts.user_id") {
-		t.Error("Output should contain posts.user_id in reference")
-	}
-	if !strings.Contains(output, "public.users.id") {
-		t.Error("Output should contain users.id in reference")
-	}
-	if !strings.Contains(output, "ondelete: CASCADE") {
-		t.Error("Output should contain ondelete: CASCADE")
-	}
-	if !strings.Contains(output, "onupdate: CASCADE") {
-		t.Error("Output should contain onupdate: CASCADE")
-	}
-
-	// Verify index
-	if !strings.Contains(output, "indexes") {
-		t.Error("Output should contain indexes section")
-	}
-	if !strings.Contains(output, "(email)") {
-		t.Error("Output should contain email index")
-	}
-	if !strings.Contains(output, "unique") {
-		t.Error("Output should contain unique attribute for email index")
-	}
+	assert.Contains(t, output, "Table public.users {")
+	assert.Contains(t, output, "Table public.posts {")
+	assert.Contains(t, output, "Ref: public.posts.user_id > public.users.id [delete: CASCADE]")
+	assert.Contains(t, output, "(email) [unique, name: 'idx_users_email']")
 }

-func TestWriter_WriteSchema(t *testing.T) {
+func TestWriter_WriteDatabase_OneToOneRelationship(t *testing.T) {
+	db := models.InitDatabase("test_db")
 	schema := models.InitSchema("public")

-	table := models.InitTable("users", "public")
-	idCol := models.InitColumn("id", "users", "public")
-	idCol.Type = "bigint"
-	idCol.IsPrimaryKey = true
-	idCol.NotNull = true
-	table.Columns["id"] = idCol
-
-	usernameCol := models.InitColumn("username", "users", "public")
-	usernameCol.Type = "varchar(50)"
-	usernameCol.NotNull = true
-	table.Columns["username"] = usernameCol
-
-	schema.Tables = append(schema.Tables, table)
-
-	tmpDir := t.TempDir()
-	outputPath := filepath.Join(tmpDir, "test.dbml")
-
-	opts := &writers.WriterOptions{
-		OutputPath: outputPath,
-	}
-
-	writer := NewWriter(opts)
-	err := writer.WriteSchema(schema)
-	if err != nil {
-		t.Fatalf("WriteSchema() error = %v", err)
-	}
-
-	content, err := os.ReadFile(outputPath)
-	if err != nil {
-		t.Fatalf("Failed to read output file: %v", err)
-	}
-
-	output := string(content)
-
-	// Verify table exists
-	if !strings.Contains(output, "Table public.users {") {
-		t.Error("Output should contain users table")
-	}
-
-	// Verify columns
-	if !strings.Contains(output, "id bigint") {
-		t.Error("Output should contain id column")
-	}
-	if !strings.Contains(output, "username varchar(50)") {
-		t.Error("Output should contain username column")
-	}
-}
-
-func TestWriter_WriteDatabase_MultipleSchemas(t *testing.T) {
-	db := models.InitDatabase("test_db")
-
-	// Create public schema with users table
-	publicSchema := models.InitSchema("public")
 	usersTable := models.InitTable("users", "public")
 	idCol := models.InitColumn("id", "users", "public")
 	idCol.Type = "bigint"
 	idCol.IsPrimaryKey = true
 	usersTable.Columns["id"] = idCol
-	publicSchema.Tables = append(publicSchema.Tables, usersTable)
+	schema.Tables = append(schema.Tables, usersTable)

-	// Create admin schema with audit_logs table
-	adminSchema := models.InitSchema("admin")
-	auditTable := models.InitTable("audit_logs", "admin")
-	auditIdCol := models.InitColumn("id", "audit_logs", "admin")
-	auditIdCol.Type = "bigint"
-	auditIdCol.IsPrimaryKey = true
-	auditTable.Columns["id"] = auditIdCol
-
-	userIdCol := models.InitColumn("user_id", "audit_logs", "admin")
+	profilesTable := models.InitTable("profiles", "public")
+	profileIdCol := models.InitColumn("id", "profiles", "public")
+	profileIdCol.Type = "bigint"
+	profilesTable.Columns["id"] = profileIdCol
+	userIdCol := models.InitColumn("user_id", "profiles", "public")
 	userIdCol.Type = "bigint"
-	auditTable.Columns["user_id"] = userIdCol
+	userIdCol.IsPrimaryKey = true // This makes it a one-to-one
+	profilesTable.Columns["user_id"] = userIdCol

-	// Add foreign key from admin.audit_logs to public.users
-	fk := models.InitConstraint("fk_audit_user", models.ForeignKeyConstraint)
-	fk.Table = "audit_logs"
-	fk.Schema = "admin"
+	fk := models.InitConstraint("fk_profiles_user", models.ForeignKeyConstraint)
+	fk.Table = "profiles"
+	fk.Schema = "public"
 	fk.Columns = []string{"user_id"}
 	fk.ReferencedTable = "users"
 	fk.ReferencedSchema = "public"
 	fk.ReferencedColumns = []string{"id"}
-	fk.OnDelete = "SET NULL"
-	auditTable.Constraints["fk_audit_user"] = fk
-
-	adminSchema.Tables = append(adminSchema.Tables, auditTable)
-
-	db.Schemas = append(db.Schemas, publicSchema, adminSchema)
-
-	tmpDir := t.TempDir()
-	outputPath := filepath.Join(tmpDir, "test.dbml")
-
-	opts := &writers.WriterOptions{
-		OutputPath: outputPath,
-	}
-
-	writer := NewWriter(opts)
-	err := writer.WriteDatabase(db)
-	if err != nil {
-		t.Fatalf("WriteDatabase() error = %v", err)
-	}
-
-	content, err := os.ReadFile(outputPath)
-	if err != nil {
-		t.Fatalf("Failed to read output file: %v", err)
-	}
-
-	output := string(content)
-
-	// Verify both schemas present
-	if !strings.Contains(output, "public.users") {
-		t.Error("Output should contain public.users table")
-	}
-	if !strings.Contains(output, "admin.audit_logs") {
-		t.Error("Output should contain admin.audit_logs table")
-	}
-
-	// Verify cross-schema foreign key
-	if !strings.Contains(output, "admin.audit_logs.user_id") {
-		t.Error("Output should contain admin.audit_logs.user_id in reference")
-	}
-	if !strings.Contains(output, "public.users.id") {
-		t.Error("Output should contain public.users.id in reference")
-	}
-	if !strings.Contains(output, "ondelete: SET NULL") {
-		t.Error("Output should contain ondelete: SET NULL")
-	}
-}
-
-func TestWriter_WriteTable_WithDefaults(t *testing.T) {
-	table := models.InitTable("products", "public")
-
-	idCol := models.InitColumn("id", "products", "public")
-	idCol.Type = "bigint"
-	idCol.IsPrimaryKey = true
-	table.Columns["id"] = idCol
-
-	isActiveCol := models.InitColumn("is_active", "products", "public")
-	isActiveCol.Type = "boolean"
-	isActiveCol.Default = "true"
-	table.Columns["is_active"] = isActiveCol
-
-	createdCol := models.InitColumn("created_at", "products", "public")
-	createdCol.Type = "timestamp"
-	createdCol.Default = "CURRENT_TIMESTAMP"
-	table.Columns["created_at"] = createdCol
-
-	tmpDir := t.TempDir()
-	outputPath := filepath.Join(tmpDir, "test.dbml")
-
-	opts := &writers.WriterOptions{
-		OutputPath: outputPath,
-	}
-
-	writer := NewWriter(opts)
-	err := writer.WriteTable(table)
-	if err != nil {
-		t.Fatalf("WriteTable() error = %v", err)
-	}
-
-	content, err := os.ReadFile(outputPath)
-	if err != nil {
-		t.Fatalf("Failed to read output file: %v", err)
-	}
-
-	output := string(content)
-
-	// Verify default values
-	if !strings.Contains(output, "default:") {
-		t.Error("Output should contain default values")
-	}
-}
-
-func TestWriter_WriteTable_EmptyPath(t *testing.T) {
-	table := models.InitTable("users", "public")
-	idCol := models.InitColumn("id", "users", "public")
-	idCol.Type = "bigint"
-	table.Columns["id"] = idCol
-
-	// When OutputPath is empty, it should print to stdout (not error)
-	opts := &writers.WriterOptions{
-		OutputPath: "",
-	}
-
-	writer := NewWriter(opts)
-	err := writer.WriteTable(table)
-	if err != nil {
-		t.Fatalf("WriteTable() with empty path should not error, got: %v", err)
-	}
-}
-
-func TestWriter_WriteDatabase_WithComments(t *testing.T) {
-	db := models.InitDatabase("test_db")
-	db.Description = "Test database description"
-	db.Comment = "Additional comment"
-
-	schema := models.InitSchema("public")
-	table := models.InitTable("users", "public")
-	table.Comment = "Users table comment"
-
-	idCol := models.InitColumn("id", "users", "public")
-	idCol.Type = "bigint"
-	idCol.IsPrimaryKey = true
-	idCol.Comment = "Primary key"
-	table.Columns["id"] = idCol
-
-	schema.Tables = append(schema.Tables, table)
+	profilesTable.Constraints["fk_profiles_user"] = fk
+	schema.Tables = append(schema.Tables, profilesTable)
 	db.Schemas = append(db.Schemas, schema)

 	tmpDir := t.TempDir()
 	outputPath := filepath.Join(tmpDir, "test.dbml")
-
-	opts := &writers.WriterOptions{
-		OutputPath: outputPath,
-	}
-
+	opts := &writers.WriterOptions{OutputPath: outputPath}
 	writer := NewWriter(opts)
 	err := writer.WriteDatabase(db)
-	if err != nil {
-		t.Fatalf("WriteDatabase() error = %v", err)
-	}
+	assert.NoError(t, err)

 	content, err := os.ReadFile(outputPath)
-	if err != nil {
-		t.Fatalf("Failed to read output file: %v", err)
-	}
-
+	assert.NoError(t, err)
 	output := string(content)

-	// Verify comments are present
-	if !strings.Contains(output, "//") {
-		t.Error("Output should contain comments")
-	}
-}
-
-func TestWriter_WriteDatabase_WithIndexType(t *testing.T) {
-	db := models.InitDatabase("test_db")
-	schema := models.InitSchema("public")
-	table := models.InitTable("users", "public")
-
-	idCol := models.InitColumn("id", "users", "public")
-	idCol.Type = "bigint"
-	idCol.IsPrimaryKey = true
-	table.Columns["id"] = idCol
-
-	emailCol := models.InitColumn("email", "users", "public")
-	emailCol.Type = "varchar(255)"
-	table.Columns["email"] = emailCol
-
-	// Add index with type
-	idx := models.InitIndex("idx_email")
-	idx.Columns = []string{"email"}
-	idx.Type = "btree"
-	idx.Unique = true
-	idx.Table = "users"
-	idx.Schema = "public"
-	table.Indexes["idx_email"] = idx
-
-	schema.Tables = append(schema.Tables, table)
-	db.Schemas = append(db.Schemas, schema)
-
-	tmpDir := t.TempDir()
-	outputPath := filepath.Join(tmpDir, "test.dbml")
-
-	opts := &writers.WriterOptions{
-		OutputPath: outputPath,
-	}
-
-	writer := NewWriter(opts)
-	err := writer.WriteDatabase(db)
-	if err != nil {
-		t.Fatalf("WriteDatabase() error = %v", err)
-	}
-
-	content, err := os.ReadFile(outputPath)
-	if err != nil {
-		t.Fatalf("Failed to read output file: %v", err)
-	}
-
-	output := string(content)
-
-	// Verify index with type
-	if !strings.Contains(output, "type:") || !strings.Contains(output, "btree") {
-		t.Error("Output should contain index type")
-	}
-}
+	assert.Contains(t, output, "Ref: public.profiles.user_id - public.users.id")
+}