github.com/iasthc/atlas/cmd/atlas@v0.0.0-20230523071841-73246df3f88d/internal/migrate/ent/enttest/enttest.go (about)

     1  // Copyright 2021-present The Atlas Authors. All rights reserved.
     2  // This source code is licensed under the Apache 2.0 license found
     3  // in the LICENSE file in the root directory of this source tree.
     4  
     5  // Code generated by entc, DO NOT EDIT.
     6  
     7  package enttest
     8  
     9  import (
    10  	"context"
    11  
    12  	"github.com/iasthc/atlas/cmd/atlas/internal/migrate/ent"
    13  	// required by schema hooks.
    14  	_ "github.com/iasthc/atlas/cmd/atlas/internal/migrate/ent/runtime"
    15  
    16  	"github.com/iasthc/atlas/cmd/atlas/internal/migrate/ent/migrate"
    17  	"entgo.io/ent/dialect/sql/schema"
    18  )
    19  
    20  type (
    21  	// TestingT is the interface that is shared between
    22  	// testing.T and testing.B and used by enttest.
    23  	TestingT interface {
    24  		FailNow()
    25  		Error(...any)
    26  	}
    27  
    28  	// Option configures client creation.
    29  	Option func(*options)
    30  
    31  	options struct {
    32  		opts        []ent.Option
    33  		migrateOpts []schema.MigrateOption
    34  	}
    35  )
    36  
    37  // WithOptions forwards options to client creation.
    38  func WithOptions(opts ...ent.Option) Option {
    39  	return func(o *options) {
    40  		o.opts = append(o.opts, opts...)
    41  	}
    42  }
    43  
    44  // WithMigrateOptions forwards options to auto migration.
    45  func WithMigrateOptions(opts ...schema.MigrateOption) Option {
    46  	return func(o *options) {
    47  		o.migrateOpts = append(o.migrateOpts, opts...)
    48  	}
    49  }
    50  
    51  func newOptions(opts []Option) *options {
    52  	o := &options{}
    53  	for _, opt := range opts {
    54  		opt(o)
    55  	}
    56  	return o
    57  }
    58  
    59  // Open calls ent.Open and auto-run migration.
    60  func Open(t TestingT, driverName, dataSourceName string, opts ...Option) *ent.Client {
    61  	o := newOptions(opts)
    62  	c, err := ent.Open(driverName, dataSourceName, o.opts...)
    63  	if err != nil {
    64  		t.Error(err)
    65  		t.FailNow()
    66  	}
    67  	migrateSchema(t, c, o)
    68  	return c
    69  }
    70  
    71  // NewClient calls ent.NewClient and auto-run migration.
    72  func NewClient(t TestingT, opts ...Option) *ent.Client {
    73  	o := newOptions(opts)
    74  	c := ent.NewClient(o.opts...)
    75  	migrateSchema(t, c, o)
    76  	return c
    77  }
    78  func migrateSchema(t TestingT, c *ent.Client, o *options) {
    79  	tables, err := schema.CopyTables(migrate.Tables)
    80  	if err != nil {
    81  		t.Error(err)
    82  		t.FailNow()
    83  	}
    84  	if err := migrate.Create(context.Background(), c.Schema, tables, o.migrateOpts...); err != nil {
    85  		t.Error(err)
    86  		t.FailNow()
    87  	}
    88  }