github.com/ledgerwatch/erigon-lib@v1.0.0/kv/memdb/memory_database.go (about)

     1  /*
     2     Copyright 2021 Erigon contributors
     3  
     4     Licensed under the Apache License, Version 2.0 (the "License");
     5     you may not use this file except in compliance with the License.
     6     You may obtain a copy of the License at
     7  
     8         http://www.apache.org/licenses/LICENSE-2.0
     9  
    10     Unless required by applicable law or agreed to in writing, software
    11     distributed under the License is distributed on an "AS IS" BASIS,
    12     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13     See the License for the specific language governing permissions and
    14     limitations under the License.
    15  */
    16  
    17  package memdb
    18  
    19  import (
    20  	"context"
    21  	"testing"
    22  
    23  	"github.com/ledgerwatch/erigon-lib/kv"
    24  	"github.com/ledgerwatch/erigon-lib/kv/mdbx"
    25  	"github.com/ledgerwatch/log/v3"
    26  )
    27  
    28  func New(tmpDir string) kv.RwDB {
    29  	return mdbx.NewMDBX(log.New()).InMem(tmpDir).MustOpen()
    30  }
    31  
    32  func NewPoolDB(tmpDir string) kv.RwDB {
    33  	return mdbx.NewMDBX(log.New()).InMem(tmpDir).Label(kv.TxPoolDB).WithTableCfg(func(_ kv.TableCfg) kv.TableCfg { return kv.TxpoolTablesCfg }).MustOpen()
    34  }
    35  func NewDownloaderDB(tmpDir string) kv.RwDB {
    36  	return mdbx.NewMDBX(log.New()).InMem(tmpDir).Label(kv.DownloaderDB).WithTableCfg(func(_ kv.TableCfg) kv.TableCfg { return kv.DownloaderTablesCfg }).MustOpen()
    37  }
    38  func NewSentryDB(tmpDir string) kv.RwDB {
    39  	return mdbx.NewMDBX(log.New()).InMem(tmpDir).Label(kv.SentryDB).WithTableCfg(func(_ kv.TableCfg) kv.TableCfg { return kv.SentryTablesCfg }).MustOpen()
    40  }
    41  
    42  func NewTestDB(tb testing.TB) kv.RwDB {
    43  	tb.Helper()
    44  	tmpDir := tb.TempDir()
    45  	tb.Helper()
    46  	db := New(tmpDir)
    47  	tb.Cleanup(db.Close)
    48  	return db
    49  }
    50  
    51  func BeginRw(tb testing.TB, db kv.RwDB) kv.RwTx {
    52  	tb.Helper()
    53  	tx, err := db.BeginRw(context.Background())
    54  	if err != nil {
    55  		tb.Fatal(err)
    56  	}
    57  	tb.Cleanup(tx.Rollback)
    58  	return tx
    59  }
    60  
    61  func BeginRo(tb testing.TB, db kv.RoDB) kv.Tx {
    62  	tb.Helper()
    63  	tx, err := db.BeginRo(context.Background())
    64  	if err != nil {
    65  		tb.Fatal(err)
    66  	}
    67  	tb.Cleanup(tx.Rollback)
    68  	return tx
    69  }
    70  
    71  func NewTestPoolDB(tb testing.TB) kv.RwDB {
    72  	tb.Helper()
    73  	tmpDir := tb.TempDir()
    74  	db := NewPoolDB(tmpDir)
    75  	tb.Cleanup(db.Close)
    76  	return db
    77  }
    78  
    79  func NewTestDownloaderDB(tb testing.TB) kv.RwDB {
    80  	tb.Helper()
    81  	tmpDir := tb.TempDir()
    82  	db := NewDownloaderDB(tmpDir)
    83  	tb.Cleanup(db.Close)
    84  	return db
    85  }
    86  
    87  func NewTestSentrylDB(tb testing.TB) kv.RwDB {
    88  	tb.Helper()
    89  	tmpDir := tb.TempDir()
    90  	db := NewPoolDB(tmpDir)
    91  	tb.Cleanup(db.Close)
    92  	return db
    93  }
    94  
    95  func NewTestTx(tb testing.TB) (kv.RwDB, kv.RwTx) {
    96  	tb.Helper()
    97  	tmpDir := tb.TempDir()
    98  	db := New(tmpDir)
    99  	tb.Cleanup(db.Close)
   100  	tx, err := db.BeginRw(context.Background())
   101  	if err != nil {
   102  		tb.Fatal(err)
   103  	}
   104  	tb.Cleanup(tx.Rollback)
   105  	return db, tx
   106  }
   107  
   108  func NewTestPoolTx(tb testing.TB) (kv.RwDB, kv.RwTx) {
   109  	tb.Helper()
   110  	db := NewTestPoolDB(tb)
   111  	tx, err := db.BeginRw(context.Background()) //nolint
   112  	if err != nil {
   113  		tb.Fatal(err)
   114  	}
   115  	if tb != nil {
   116  		tb.Cleanup(tx.Rollback)
   117  	}
   118  	return db, tx
   119  }
   120  
   121  func NewTestSentryTx(tb testing.TB) (kv.RwDB, kv.RwTx) {
   122  	tb.Helper()
   123  	db := NewTestSentrylDB(tb)
   124  	tx, err := db.BeginRw(context.Background()) //nolint
   125  	if err != nil {
   126  		tb.Fatal(err)
   127  	}
   128  	if tb != nil {
   129  		tb.Cleanup(tx.Rollback)
   130  	}
   131  	return db, tx
   132  }