From f59d4dc5b81898de12d676012f88db377399d6b6 Mon Sep 17 00:00:00 2001
From: Knut Ahlers <knut@ahlers.me>
Date: Sun, 25 Jul 2021 13:42:56 +0200
Subject: [PATCH] Add support for "Expect" check, add Duration support

Signed-off-by: Knut Ahlers <knut@ahlers.me>
---
 attributeStore.go | 42 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/attributeStore.go b/attributeStore.go
index f8279f1..ddc7414 100644
--- a/attributeStore.go
+++ b/attributeStore.go
@@ -1,8 +1,11 @@
 package main
 
 import (
-	"errors"
 	"fmt"
+	"strings"
+	"time"
+
+	"github.com/pkg/errors"
 )
 
 var (
@@ -12,6 +15,33 @@ var (
 
 type moduleAttributeStore map[string]interface{}
 
+func (m moduleAttributeStore) Expect(keys ...string) error {
+	var missing []string
+
+	for _, k := range keys {
+		if _, ok := m[k]; !ok {
+			missing = append(missing, k)
+		}
+	}
+
+	if len(missing) > 0 {
+		return errors.Errorf("missing key(s) %s", strings.Join(missing, ", "))
+	}
+
+	return nil
+}
+
+func (m moduleAttributeStore) MustDuration(name string, defVal *time.Duration) time.Duration {
+	v, err := m.Duration(name)
+	if err != nil {
+		if defVal != nil {
+			return *defVal
+		}
+		panic(err)
+	}
+	return v
+}
+
 func (m moduleAttributeStore) MustInt64(name string, defVal *int64) int64 {
 	v, err := m.Int64(name)
 	if err != nil {
@@ -34,6 +64,16 @@ func (m moduleAttributeStore) MustString(name string, defVal *string) string {
 	return v
 }
 
+func (m moduleAttributeStore) Duration(name string) (time.Duration, error) {
+	v, err := m.String(name)
+	if err != nil {
+		return 0, errors.Wrap(err, "getting string value")
+	}
+
+	d, err := time.ParseDuration(v)
+	return d, errors.Wrap(err, "parsing value")
+}
+
 func (m moduleAttributeStore) Int64(name string) (int64, error) {
 	v, ok := m[name]
 	if !ok {