2 files changed, 110 insertions, 37 deletions

diff --git a/pkg/assert/assert.go b/pkg/assert/assert.go
index 7281a95..65fe088 100644
--- a/pkg/assert/assert.go
+++ b/pkg/assert/assert.go
@@ -1,25 +1,94 @@
 package assert
 
-import "log"
+import (
+	"fmt"
+	"io"
+	"os"
+	"reflect"
+	"runtime/debug"
+	"sync"
+)
 
-func Assert(assertion bool, message string, a ...any) {
+var (
+	writer io.Writer = os.Stderr
+
+	flushes = []io.Closer{}
+	flushMu sync.Mutex
+
+	assertData = map[string]any{}
+	mapMu      sync.Mutex
+)
+
+func AddData(key string, value any) {
+	mapMu.Lock()
+	assertData[key] = value
+	mapMu.Unlock()
+}
+
+func RemoveData(key string) {
+	mapMu.Lock()
+	delete(assertData, key)
+	mapMu.Unlock()
+}
+
+func AddFlush(flusher io.Closer) {
+	flushMu.Lock()
+	flushes = append(flushes, flusher)
+	flushMu.Unlock()
+}
+
+func SetWriter(w io.Writer) {
+	writer = w
+}
+
+func runAssert(message string, args ...any) {
+	flushMu.Lock()
+	for len(flushes) != 0 {
+		flusher := flushes[len(flushes)-1]
+		flusher.Close()
+		flushes = flushes[:len(flushes)-1]
+	}
+	flushMu.Unlock()
+
+	values := []any{
+		"msg", message,
+	}
+	values = append(values, args...)
+	mapMu.Lock()
+	for k, v := range assertData {
+		values = append(values, k, v)
+	}
+	mapMu.Unlock()
+
+	fmt.Fprintf(writer, "ARGS: %+v\n", args)
+	fmt.Fprintf(writer, "ASSERT\n")
+	for i := 0; i < len(values); i += 2 {
+		fmt.Fprintf(writer, "   %s=%v\n", values[i], values[i+1])
+	}
+	fmt.Fprintln(writer, string(debug.Stack()))
+
+	os.Exit(1)
+}
+
+func Assert(assertion bool, message string, args ...any) {
 	if !assertion {
-		log.Fatalf(message+"\n", a...)
+		runAssert(message, args...)
 	}
 }
 
-func NoError(err error, message string, a ...any) {
+func NoError(err error, message string, args ...any) {
 	if err != nil {
-		log.Fatalf(message+": "+err.Error()+"\n", a...)
+		args = append(args, "error", err)
+		runAssert(message, args...)
 	}
 }
 
-func Never(message string, a ...any) {
-	log.Fatalf(message+"\n", a...)
+func Never(message string, args ...any) {
+	runAssert(message, args...)
 }
 
-func NotNil(value any, message string, a ...any) {
-	if value == nil {
-		log.Fatalf(message+"\n", a...)
+func NotNil(value any, message string, args ...any) {
+	if value == nil || reflect.ValueOf(value).Kind() == reflect.Ptr && reflect.ValueOf(value).IsNil() {
+		runAssert(message, args...)
 	}
 }
diff --git a/pkg/snowflake/snowflake.go b/pkg/snowflake/snowflake.go
index bb4a07b..ab5fafa 100644
--- a/pkg/snowflake/snowflake.go
+++ b/pkg/snowflake/snowflake.go
@@ -12,16 +12,34 @@ import (
 const (
 	// Epoch is set to the twitter snowflake epoch of Nov 04 2010 01:42:54 UTC in milliseconds
 	// TODO: change this to eko epoch when eko is production ready
-	Epoch int64 = 1288834974657
-	nodeBits  = 10
-	stepBits  = 12
-	NodeMax   = 1<<nodeBits - 1
-	nodeMask  = NodeMax << stepBits
-	stepMask  = 1<<stepBits - 1
-	timeShift = nodeBits + stepBits
-	nodeShift = stepBits
+	Epoch     int64 = 1288834974657
+	nodeBits        = 10
+	stepBits        = 12
+	NodeMax         = 1<<nodeBits - 1
+	nodeMask        = NodeMax << stepBits
+	stepMask        = 1<<stepBits - 1
+	timeShift       = nodeBits + stepBits
+	nodeShift       = stepBits
 )
 
+type ID int64
+
+func (id ID) String() string {
+	return strconv.FormatInt(int64(id), 10)
+}
+
+func (id ID) Time() (msSinceEpoch int64) {
+	return (int64(id) >> timeShift) + Epoch
+}
+
+func (id ID) Node() int64 {
+	return int64(id) & nodeMask >> nodeShift
+}
+
+func (id ID) Step() int64 {
+	return int64(id) & stepMask
+}
+
 type Node struct {
 	mu    sync.Mutex
 	epoch time.Time
@@ -30,11 +48,13 @@ type Node struct {
 	step  int64
 }
 
-type ID int64
+func (n *Node) String() string {
+	return fmt.Sprintf("Node%v(epoch: %v, step: %v, time: %v)", n.node, n.epoch, n.step, n.time)
+}
 
 func NewNode(node int64) *Node {
-	assert.Assert(nodeBits+stepBits <= 22, "node and step bits must add up to 22 or less")
-	assert.Assert(0 <= node && node <= NodeMax, "node and step bits must add up to 22 or less")
+	assert.Assert(nodeBits+stepBits == 22, "node and step bits must add up to 22")
+	assert.Assert(0 <= node && node <= NodeMax, "node must be within 0 and NodeMax", "node", node)
 
 	// Credit to https://github.com/bwmarrin/snowflake
 	currentTime := time.Now()
@@ -67,22 +87,6 @@ func (n *Node) Generate() ID {
 	return ID((now << timeShift) | (n.node << nodeShift) | (n.step))
 }
 
-func (id ID) String() string {
-	return strconv.FormatInt(int64(id), 10)
-}
-
-func (id ID) Time() (msSinceEpoch int64) {
-	return (int64(id) >> timeShift) + Epoch
-}
-
-func (id ID) Node() int64 {
-	return int64(id) & nodeMask >> nodeShift
-}
-
-func (id ID) Step() int64 {
-	return int64(id) & stepMask
-}
-
 // Json marshling to avoid inprecision of json number (float64)
 
 func (id ID) MarshalJSON() ([]byte, error) {