From 5c401f9c06d65696ccc05069921197d6454ae718 Mon Sep 17 00:00:00 2001
From: Yoshiyuki Mineo <Yoshiyuki.Mineo@jp.sony.com>
Date: Sat, 28 Jun 2025 19:24:22 +0900
Subject: [PATCH] Make unit tests stabler (#78)

* Refactor Sonyflake to use a customizable time function for improved testability. Update currentElapsedTime and sleep methods to utilize the new time function instead of directly calling time.Now().

* Update TimeUnit in ToTime test to use time.Millisecond for improved accuracy in timestamp validation. Adjust time duration comparison to ensure correct validation of generated timestamps.

* Refactor TestNextID to use a customizable time function for improved accuracy in timestamp validation. Adjust time comparison to ensure expected results in generated IDs.

* Refactor TestNextID_InSequence to use a consistent start time for improved clarity in timestamp validation. Update max sequence comparison to ensure accurate validation of generated IDs.

* Refactor tests in sonyflake_test.go to replace fmt.Println with t.Log for better test output management. Update error handling to use errors.New for consistency.
---
 v2/sonyflake.go      |  7 +++++--
 v2/sonyflake_test.go | 37 +++++++++++++++++++------------------
 2 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/v2/sonyflake.go b/v2/sonyflake.go
index 00bf609..0ca2be4 100644
--- a/v2/sonyflake.go
+++ b/v2/sonyflake.go
@@ -67,6 +67,8 @@ type Sonyflake struct {
 
 	sequence int
 	machine  int
+
+	now func() time.Time
 }
 
 var (
@@ -116,6 +118,7 @@ func New(st Settings) (*Sonyflake, error) {
 
 	sf := new(Sonyflake)
 	sf.mutex = new(sync.Mutex)
+	sf.now = time.Now
 
 	if st.BitsSequence == 0 {
 		sf.bitsSequence = defaultBitsSequence
@@ -198,12 +201,12 @@ func (sf *Sonyflake) toInternalTime(t time.Time) int64 {
 }
 
 func (sf *Sonyflake) currentElapsedTime() int64 {
-	return sf.toInternalTime(time.Now()) - sf.startTime
+	return sf.toInternalTime(sf.now()) - sf.startTime
 }
 
 func (sf *Sonyflake) sleep(overtime int64) {
 	sleepTime := time.Duration(overtime*sf.timeUnit) -
-		time.Duration(time.Now().UTC().UnixNano()%sf.timeUnit)
+		time.Duration(sf.now().UTC().UnixNano()%sf.timeUnit)
 	time.Sleep(sleepTime)
 }
 
diff --git a/v2/sonyflake_test.go b/v2/sonyflake_test.go
index 91fc060..6a95758 100644
--- a/v2/sonyflake_test.go
+++ b/v2/sonyflake_test.go
@@ -2,7 +2,6 @@ package sonyflake
 
 import (
 	"errors"
-	"fmt"
 	"net"
 	"runtime"
 	"testing"
@@ -13,7 +12,7 @@ import (
 )
 
 func TestNew(t *testing.T) {
-	errGetMachineID := fmt.Errorf("failed to get machine id")
+	errGetMachineID := errors.New("failed to get machine id")
 
 	testCases := []struct {
 		name     string
@@ -138,15 +137,16 @@ func defaultMachineID(t *testing.T) int {
 }
 
 func TestNextID(t *testing.T) {
-	sf := newSonyflake(t, Settings{StartTime: time.Now()})
+	start := time.Now()
+	sf := newSonyflake(t, Settings{StartTime: start})
 
 	sleepTime := int64(50)
-	time.Sleep(time.Duration(sleepTime * sf.timeUnit))
+	sf.now = func() time.Time { return start.Add(time.Duration(sleepTime * sf.timeUnit)) }
 
 	id := nextID(t, sf)
 
 	actualTime := sf.timePart(id)
-	if actualTime < sleepTime || actualTime > sleepTime+1 {
+	if actualTime != sleepTime {
 		t.Errorf("unexpected time: %d", actualTime)
 	}
 
@@ -160,17 +160,17 @@ func TestNextID(t *testing.T) {
 		t.Errorf("unexpected machine: %d", actualMachine)
 	}
 
-	fmt.Println("sonyflake id:", id)
-	fmt.Println("decompose:", sf.Decompose(id))
+	t.Log("sonyflake id:", id)
+	t.Log("decompose:", sf.Decompose(id))
 }
 
 func TestNextID_InSequence(t *testing.T) {
-	now := time.Now()
+	start := time.Now()
 	sf := newSonyflake(t, Settings{
 		TimeUnit:  time.Millisecond,
-		StartTime: now,
+		StartTime: start,
 	})
-	startTime := sf.toInternalTime(now)
+	startTime := sf.toInternalTime(start)
 	machineID := int64(defaultMachineID(t))
 
 	var numID int
@@ -210,11 +210,11 @@ func TestNextID_InSequence(t *testing.T) {
 		}
 	}
 
-	if maxSeq != 1<<sf.bitsSequence-1 {
+	if maxSeq > 1<<sf.bitsSequence-1 {
 		t.Errorf("unexpected max sequence: %d", maxSeq)
 	}
-	fmt.Println("max sequence:", maxSeq)
-	fmt.Println("number of id:", numID)
+	t.Log("max sequence:", maxSeq)
+	t.Log("number of id:", numID)
 }
 
 func TestNextID_InParallel(t *testing.T) {
@@ -223,7 +223,7 @@ func TestNextID_InParallel(t *testing.T) {
 
 	numCPU := runtime.NumCPU()
 	runtime.GOMAXPROCS(numCPU)
-	fmt.Println("number of cpu:", numCPU)
+	t.Log("number of cpu:", numCPU)
 
 	consumer := make(chan int64)
 
@@ -250,7 +250,7 @@ func TestNextID_InParallel(t *testing.T) {
 		}
 		set[id] = struct{}{}
 	}
-	fmt.Println("number of id:", len(set))
+	t.Log("number of id:", len(set))
 }
 
 func pseudoSleep(sf *Sonyflake, period time.Duration) {
@@ -358,16 +358,17 @@ func TestLower16BitPrivateIP(t *testing.T) {
 func TestToTime(t *testing.T) {
 	start := time.Now()
 	sf := newSonyflake(t, Settings{
-		TimeUnit:  100 * time.Millisecond,
+		TimeUnit:  time.Millisecond,
 		StartTime: start,
 	})
 
+	sf.now = func() time.Time { return start }
 	id := nextID(t, sf)
 
 	tm := sf.ToTime(id)
 	diff := tm.Sub(start)
-	if diff < 0 || diff > time.Duration(sf.timeUnit) {
-		t.Errorf("unexpected time: %v", tm)
+	if diff < 0 || diff >= time.Duration(sf.timeUnit) {
+		t.Errorf("unexpected time: %v", diff)
 	}
 }