From c4d2d7c9b8263c76ebeb334541b821d2fa758579 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Robert=20Paj=C4=85k?= <pellared@hotmail.com>
Date: Wed, 2 Aug 2023 16:52:04 +0200
Subject: [PATCH] sdk/metric: Add unit tests for Shutdown WithTimeout (#4379)

* sdk/metric: Improve WithTimeout doc

* Add tests
---
 sdk/metric/periodic_reader_test.go | 46 ++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/sdk/metric/periodic_reader_test.go b/sdk/metric/periodic_reader_test.go
index e5b279ec238..6ac527a37ea 100644
--- a/sdk/metric/periodic_reader_test.go
+++ b/sdk/metric/periodic_reader_test.go
@@ -329,6 +329,52 @@ func TestPeriodicReaderFlushesPending(t *testing.T) {
 		_ = r.Shutdown(context.Background())
 	})
 
+	t.Run("ForceFlush timeout on producer", func(t *testing.T) {
+		exp, called := expFunc(t)
+		timeout := time.Millisecond
+		r := NewPeriodicReader(exp, WithTimeout(timeout))
+		r.register(testSDKProducer{
+			produceFunc: func(ctx context.Context, rm *metricdata.ResourceMetrics) error {
+				select {
+				case <-time.After(timeout + time.Second):
+					*rm = testResourceMetricsA
+				case <-ctx.Done():
+					// we timed out before we could collect metrics
+					return ctx.Err()
+				}
+				return nil
+			}})
+		r.RegisterProducer(testExternalProducer{})
+		assert.Equal(t, context.DeadlineExceeded, r.ForceFlush(context.Background()), "timeout error not returned")
+		assert.False(t, *called, "exporter Export method called when it should have failed before export")
+
+		// Ensure Reader is allowed clean up attempt.
+		_ = r.Shutdown(context.Background())
+	})
+
+	t.Run("ForceFlush timeout on external producer", func(t *testing.T) {
+		exp, called := expFunc(t)
+		timeout := time.Millisecond
+		r := NewPeriodicReader(exp, WithTimeout(timeout))
+		r.register(testSDKProducer{})
+		r.RegisterProducer(testExternalProducer{
+			produceFunc: func(ctx context.Context) ([]metricdata.ScopeMetrics, error) {
+				select {
+				case <-time.After(timeout + time.Second):
+				case <-ctx.Done():
+					// we timed out before we could collect metrics
+					return nil, ctx.Err()
+				}
+				return []metricdata.ScopeMetrics{testScopeMetricsA}, nil
+			},
+		})
+		assert.Equal(t, context.DeadlineExceeded, r.ForceFlush(context.Background()), "timeout error not returned")
+		assert.False(t, *called, "exporter Export method called when it should have failed before export")
+
+		// Ensure Reader is allowed clean up attempt.
+		_ = r.Shutdown(context.Background())
+	})
+
 	t.Run("Shutdown", func(t *testing.T) {
 		exp, called := expFunc(t)
 		r := NewPeriodicReader(exp)