From 27130c410b547a48f3fb04564ff54327e9962474 Mon Sep 17 00:00:00 2001
From: Gurjot <gurjotsyd@gmail.com>
Date: Tue, 10 Dec 2024 17:01:52 +0530
Subject: [PATCH] fix

---
 cmd/staking-expiry-checker/main.go | 19 +++++-------------
 internal/poller/poller.go          | 32 ++++++++++++++++++++++++------
 2 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/cmd/staking-expiry-checker/main.go b/cmd/staking-expiry-checker/main.go
index 4ea4b74..b9dc149 100644
--- a/cmd/staking-expiry-checker/main.go
+++ b/cmd/staking-expiry-checker/main.go
@@ -65,25 +65,16 @@ func main() {
 		log.Fatal().Err(err).Msg("error while creating service")
 	}
 
-	// Create expiry poller
-	expiryPoller, err := poller.NewPoller(
-		poller.ExpiryPoller,
+	// Even though we pass service, it's viewed only through the specific interface
+	expiryPoller := poller.NewExpiryPoller(
 		cfg.Pollers.ExpiryChecker,
-		service.ProcessExpiredDelegations,
+		service, // service implements ExpiryChecker
 	)
-	if err != nil {
-		log.Fatal().Err(err).Msg("error while creating expiry poller")
-	}
 
-	// Create BTC subscriber poller
-	btcSubscriberPoller, err := poller.NewPoller(
-		poller.BTCSubscriberPoller,
+	btcSubscriberPoller := poller.NewBTCSubscriberPoller(
 		cfg.Pollers.BtcSubscriber,
-		service.ProcessBTCSubscriber,
+		service, // service implements BTCSubscriber
 	)
-	if err != nil {
-		log.Fatal().Err(err).Msg("error while creating BTC subscriber poller")
-	}
 
 	// Start pollers in separate goroutines
 	go expiryPoller.Start(ctx)
diff --git a/internal/poller/poller.go b/internal/poller/poller.go
index c54f0be..5d0f9f8 100644
--- a/internal/poller/poller.go
+++ b/internal/poller/poller.go
@@ -10,6 +10,15 @@ import (
 	"github.com/babylonlabs-io/staking-expiry-checker/internal/types"
 )
 
+// Define minimal interfaces for each poller type
+type ExpiryChecker interface {
+	ProcessExpiredDelegations(ctx context.Context) *types.Error
+}
+
+type BTCSubscriber interface {
+	ProcessBTCSubscriber(ctx context.Context) *types.Error
+}
+
 type PollerType string
 
 const (
@@ -21,20 +30,31 @@ type PollerOperation func(ctx context.Context) *types.Error
 
 type Poller struct {
 	pollerType PollerType
-	operation  PollerOperation
+	poll       func(ctx context.Context) *types.Error
 	interval   time.Duration
 	timeout    time.Duration
 	quit       chan struct{}
 }
 
-func NewPoller(pollerType PollerType, cfg config.PollerConfig, operation PollerOperation) (*Poller, error) {
+// Constructors now accept interfaces instead of the full service
+func NewExpiryPoller(cfg config.PollerConfig, checker ExpiryChecker) *Poller {
+	return &Poller{
+		pollerType: ExpiryPoller,
+		interval:   cfg.Interval,
+		timeout:    cfg.Timeout,
+		poll:       checker.ProcessExpiredDelegations,
+		quit:       make(chan struct{}),
+	}
+}
+
+func NewBTCSubscriberPoller(cfg config.PollerConfig, subscriber BTCSubscriber) *Poller {
 	return &Poller{
-		pollerType: pollerType,
-		operation:  operation,
+		pollerType: BTCSubscriberPoller,
 		interval:   cfg.Interval,
 		timeout:    cfg.Timeout,
+		poll:       subscriber.ProcessBTCSubscriber,
 		quit:       make(chan struct{}),
-	}, nil
+	}
 }
 
 func (p *Poller) Start(ctx context.Context) {
@@ -47,7 +67,7 @@ func (p *Poller) Start(ctx context.Context) {
 			pollingCtx, cancel := context.WithTimeout(ctx, p.timeout)
 			defer cancel()
 
-			if err := p.operation(pollingCtx); err != nil {
+			if err := p.poll(pollingCtx); err != nil {
 				log.Error().
 					Err(err).
 					Str("poller", string(p.pollerType)).