From 0015d25344ce91a5dc0e7ae16bdfa7757185a4bc Mon Sep 17 00:00:00 2001
From: James Rasell <jrasell@users.noreply.github.com>
Date: Wed, 19 Jul 2023 15:35:16 +0100
Subject: [PATCH] qemu: add test to cover task store functions. (#17967)

---
 drivers/qemu/state_test.go | 40 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)
 create mode 100644 drivers/qemu/state_test.go

diff --git a/drivers/qemu/state_test.go b/drivers/qemu/state_test.go
new file mode 100644
index 00000000000..be715ecd956
--- /dev/null
+++ b/drivers/qemu/state_test.go
@@ -0,0 +1,40 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: MPL-2.0
+
+package qemu
+
+import (
+	"testing"
+
+	"github.com/shoenig/test/must"
+)
+
+func Test_taskStore(t *testing.T) {
+
+	taskStoreImpl := newTaskStore()
+	must.NotNil(t, taskStoreImpl)
+
+	// Try reading something that doesn't exist.
+	taskHandleResp1, okResp1 := taskStoreImpl.Get("this-doesn't-exist")
+	must.Nil(t, taskHandleResp1)
+	must.False(t, okResp1)
+
+	// Set and get a generated task handle.
+	testTaskHandle := taskHandle{pid: 131313}
+	taskStoreImpl.Set("test-id", &testTaskHandle)
+
+	taskHandleResp2, okResp2 := taskStoreImpl.Get("test-id")
+	must.NotNil(t, taskHandleResp2)
+	must.Eq(t, &testTaskHandle, taskHandleResp2)
+	must.True(t, okResp2)
+
+	// Delete the previously set handle, and try reading it again.
+	taskStoreImpl.Delete("test-id")
+
+	taskHandleResp3, okResp3 := taskStoreImpl.Get("test-id")
+	must.Nil(t, taskHandleResp3)
+	must.False(t, okResp3)
+
+	// Deleting a non-existent handle shouldn't cause any problems.
+	taskStoreImpl.Delete("test-id")
+}