From 9f3cefecb271ea53e5f4b3fc94374272e5028dd3 Mon Sep 17 00:00:00 2001
From: josh <josh@elixxir.io>
Date: Thu, 1 Apr 2021 09:55:07 -0700
Subject: [PATCH] Integrate unknownRounds with session in storage

---
 storage/rounds/unknownRounds.go      | 11 ++++++++---
 storage/rounds/unknownRounds_test.go | 15 ++++++++++++---
 storage/session.go                   | 23 +++++++++++++++++++++++
 3 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/storage/rounds/unknownRounds.go b/storage/rounds/unknownRounds.go
index 3b0db9668..3c2b40fc1 100644
--- a/storage/rounds/unknownRounds.go
+++ b/storage/rounds/unknownRounds.go
@@ -51,23 +51,28 @@ func DefaultUnknownRoundsParams() UnknownRoundsParams {
 
 // Build and return new UnknownRounds object
 func NewUnknownRoundsStore(kv *versioned.KV,
-	params UnknownRoundsParams) *UnknownRoundsStore {
+	params UnknownRoundsParams) (*UnknownRoundsStore, error) {
 	// Build the UnmixedMessagesMap
 	// Modify the prefix of the KV
 	kv = kv.Prefix(unknownRoundPrefix)
 
-	return &UnknownRoundsStore{
+	urs:=  &UnknownRoundsStore{
 		rounds: make(map[id.Round]*uint64),
 		params: params,
 		kv:     kv,
 	}
+
+	return urs, urs.save()
 }
 
 // LoadUnknownRoundsStore loads the data for a UnknownRoundStore from disk into an object
 func LoadUnknownRoundsStore(kv *versioned.KV, params UnknownRoundsParams) (*UnknownRoundsStore, error) {
 	kv = kv.Prefix(unknownRoundPrefix)
 
-	urs := NewUnknownRoundsStore(kv, params)
+	urs, err := NewUnknownRoundsStore(kv, params)
+	if err != nil {
+		return nil, err
+	}
 
 	// Get the versioned data from the kv
 	obj, err := kv.Get(unknownRoundsStorageKey, unknownRoundsStorageVersion)
diff --git a/storage/rounds/unknownRounds_test.go b/storage/rounds/unknownRounds_test.go
index 53b9aebc9..8cc2149c4 100644
--- a/storage/rounds/unknownRounds_test.go
+++ b/storage/rounds/unknownRounds_test.go
@@ -27,7 +27,10 @@ func TestNewUnknownRoundsStore(t *testing.T) {
 		params: DefaultUnknownRoundsParams(),
 	}
 
-	store := NewUnknownRoundsStore(kv, DefaultUnknownRoundsParams())
+	store, err := NewUnknownRoundsStore(kv, DefaultUnknownRoundsParams())
+	if err != nil {
+		t.Fatalf("Failed to create known round store: %v", err)
+	}
 
 	// Compare manually created object with NewUnknownRoundsStore
 	if !reflect.DeepEqual(expectedStore, store) {
@@ -60,7 +63,10 @@ func TestNewUnknownRoundsStore(t *testing.T) {
 // Full test
 func TestUnknownRoundsStore_Iterate(t *testing.T) {
 	kv := versioned.NewKV(make(ekv.Memstore))
-	store := NewUnknownRoundsStore(kv, DefaultUnknownRoundsParams())
+	store, err := NewUnknownRoundsStore(kv, DefaultUnknownRoundsParams())
+	if err != nil {
+		t.Fatalf("Failed to create known round store: %v", err)
+	}
 
 	// Return true only for rounds that are even
 	mockChecker := func(rid id.Round) bool {
@@ -138,7 +144,10 @@ func TestUnknownRoundsStore_Iterate(t *testing.T) {
 // Unit test
 func TestLoadUnknownRoundsStore(t *testing.T) {
 	kv := versioned.NewKV(make(ekv.Memstore))
-	store := NewUnknownRoundsStore(kv, DefaultUnknownRoundsParams())
+	store, err := NewUnknownRoundsStore(kv, DefaultUnknownRoundsParams())
+	if err != nil {
+		t.Fatalf("Failed to create known round store: %v", err)
+	}
 
 	// Construct 3 lists of round IDs
 	roundListLen := 25
diff --git a/storage/session.go b/storage/session.go
index be1837de7..20c3bc903 100644
--- a/storage/session.go
+++ b/storage/session.go
@@ -10,6 +10,7 @@
 package storage
 
 import (
+	"gitlab.com/elixxir/client/storage/rounds"
 	"sync"
 	"testing"
 
@@ -62,6 +63,7 @@ type Session struct {
 	garbledMessages     *utility.MeteredCmixMessageBuffer
 	reception           *reception.Store
 	clientVersion       *clientVersion.Store
+	unknownRounds       *rounds.UnknownRoundsStore
 }
 
 // Initialize a new Session object
@@ -141,6 +143,11 @@ func New(baseDir, password string, u userInterface.User, currentVersion version.
 		return nil, errors.WithMessage(err, "Failed to create client version store.")
 	}
 
+	s.unknownRounds, err = rounds.NewUnknownRoundsStore(s.kv, rounds.DefaultUnknownRoundsParams())
+	if err != nil {
+		return nil, errors.WithMessage(err, "Failed to create unknown round store.")
+	}
+
 	return s, nil
 }
 
@@ -212,6 +219,11 @@ func Load(baseDir, password string, currentVersion version.Version,
 
 	s.reception = reception.LoadStore(s.kv)
 
+	s.unknownRounds, err = rounds.LoadUnknownRoundsStore(s.kv, rounds.DefaultUnknownRoundsParams())
+	if err != nil {
+		return nil, errors.WithMessage(err, "Failed to load unknown rounds")
+	}
+
 	return s, nil
 }
 
@@ -282,6 +294,12 @@ func (s *Session) Partition() *partition.Store {
 	return s.partition
 }
 
+func (s *Session) UnknownRounds() *rounds.UnknownRoundsStore  {
+	s.mux.RUnlock()
+	defer s.mux.RUnlock()
+	return s.unknownRounds
+}
+
 // Get an object from the session
 func (s *Session) Get(key string) (*versioned.Object, error) {
 	return s.kv.Get(key, currentSessionVersion)
@@ -363,5 +381,10 @@ func InitTestingSession(i interface{}) *Session {
 	s.partition = partition.New(s.kv)
 
 	s.reception = reception.NewStore(s.kv)
+
+	s.unknownRounds, err = rounds.NewUnknownRoundsStore(s.kv, rounds.DefaultUnknownRoundsParams())
+	if err != nil {
+		jww.FATAL.Panicf("Failed to create unknown rounds store: %+v", err)
+	}
 	return s
 }
-- 
GitLab