From 2cfc88a6fe3396b57f32d85db5cc3969fc4548f6 Mon Sep 17 00:00:00 2001
From: Matheus Nogueira <matheus.nogueira2008@gmail.com>
Date: Thu, 6 Apr 2023 16:50:00 -0300
Subject: [PATCH] fix json unmarshal and remove connection test event from
 untestable datastores (#2324)

* fix json unmarshal

* fix test

* split TraceDB and TestableTraceDB

* add TraceDB as part of TestableTraceDB
---
 server/executor/poller_executor.go     | 14 ++++++++-----
 server/http/controller.go              | 11 +++++++---
 server/model/connection_test_result.go | 28 ++++++++++++++++++++++++++
 server/tracedb/otlp.go                 |  4 ----
 server/tracedb/tracedb.go              |  6 +++++-
 5 files changed, 50 insertions(+), 13 deletions(-)

diff --git a/server/executor/poller_executor.go b/server/executor/poller_executor.go
index 73d731fc77..518776757a 100644
--- a/server/executor/poller_executor.go
+++ b/server/executor/poller_executor.go
@@ -109,11 +109,13 @@ func (pe DefaultPollerExecutor) ExecuteRequest(request *PollingRequest) (bool, s
 	}
 
 	if request.IsFirstRequest() {
-		connectionResult := traceDB.TestConnection(request.ctx)
+		if testableTraceDB, ok := traceDB.(tracedb.TestableTraceDB); ok {
+			connectionResult := testableTraceDB.TestConnection(request.ctx)
 
-		err = pe.eventEmitter.Emit(request.ctx, events.TraceDataStoreConnectionInfo(request.test.ID, request.run.ID, connectionResult))
-		if err != nil {
-			log.Printf("[PollerExecutor] Test %s Run %d: failed to emit TraceDataStoreConnectionInfo event: error: %s\n", request.test.ID, request.run.ID, err.Error())
+			err = pe.eventEmitter.Emit(request.ctx, events.TraceDataStoreConnectionInfo(request.test.ID, request.run.ID, connectionResult))
+			if err != nil {
+				log.Printf("[PollerExecutor] Test %s Run %d: failed to emit TraceDataStoreConnectionInfo event: error: %s\n", request.test.ID, request.run.ID, err.Error())
+			}
 		}
 	}
 
@@ -129,7 +131,9 @@ func (pe DefaultPollerExecutor) ExecuteRequest(request *PollingRequest) (bool, s
 
 		if !errors.Is(err, connection.ErrTraceNotFound) {
 			// run test connection to give a diagnostic when an unknown error happens
-			connectionResult = traceDB.TestConnection(request.ctx)
+			if testableTraceDB, ok := traceDB.(tracedb.TestableTraceDB); ok {
+				connectionResult = testableTraceDB.TestConnection(request.ctx)
+			}
 		}
 
 		anotherErr := pe.eventEmitter.Emit(request.ctx, events.TraceFetchingError(request.test.ID, request.run.ID, connectionResult, err))
diff --git a/server/http/controller.go b/server/http/controller.go
index bc64407e08..7f4ed05713 100644
--- a/server/http/controller.go
+++ b/server/http/controller.go
@@ -1274,10 +1274,15 @@ func (c *controller) TestConnection(ctx context.Context, dataStore openapi.DataS
 		return openapi.Response(http.StatusBadRequest, err.Error()), err
 	}
 
-	testResult := tdb.TestConnection(ctx)
+	testResult := model.ConnectionResult{}
 	statusCode := http.StatusOK
-	if !testResult.HasSucceed() {
-		statusCode = http.StatusUnprocessableEntity
+
+	if testableTraceDB, ok := tdb.(tracedb.TestableTraceDB); ok {
+		testResult = testableTraceDB.TestConnection(ctx)
+		statusCode = http.StatusOK
+		if !testResult.HasSucceed() {
+			statusCode = http.StatusUnprocessableEntity
+		}
 	}
 
 	return openapi.Response(statusCode, c.mappers.Out.ConnectionTestResult(testResult)), nil
diff --git a/server/model/connection_test_result.go b/server/model/connection_test_result.go
index 8dfece1524..3529813310 100644
--- a/server/model/connection_test_result.go
+++ b/server/model/connection_test_result.go
@@ -1,5 +1,10 @@
 package model
 
+import (
+	"encoding/json"
+	"errors"
+)
+
 type ConnectionResult struct {
 	PortCheck      ConnectionTestStep
 	Connectivity   ConnectionTestStep
@@ -18,6 +23,29 @@ type ConnectionTestStep struct {
 	Error   error
 }
 
+func (s *ConnectionTestStep) UnmarshalJSON(bytes []byte) error {
+	var step struct {
+		Passed       bool
+		Status       Status
+		Message      string
+		ErrorMessage string
+	}
+
+	err := json.Unmarshal(bytes, &step)
+	if err != nil {
+		return err
+	}
+
+	s.Passed = step.Passed
+	s.Status = step.Status
+	s.Message = step.Message
+	if step.ErrorMessage != "" {
+		s.Error = errors.New(step.ErrorMessage)
+	}
+
+	return nil
+}
+
 func (r *ConnectionTestStep) HasSucceed() bool {
 	if r == nil {
 		return true
diff --git a/server/tracedb/otlp.go b/server/tracedb/otlp.go
index c011ec14ec..ddf11bed0f 100644
--- a/server/tracedb/otlp.go
+++ b/server/tracedb/otlp.go
@@ -33,10 +33,6 @@ func (tdb *OTLPTraceDB) Close() error {
 	return nil
 }
 
-func (jtd *OTLPTraceDB) TestConnection(ctx context.Context) model.ConnectionResult {
-	return model.ConnectionResult{}
-}
-
 // GetTraceByID implements TraceDB
 func (tdb *OTLPTraceDB) GetTraceByID(ctx context.Context, id string) (model.Trace, error) {
 	run, err := tdb.db.GetRunByTraceID(ctx, traces.DecodeTraceID(id))
diff --git a/server/tracedb/tracedb.go b/server/tracedb/tracedb.go
index 445b73709b..8b4f8fcb32 100644
--- a/server/tracedb/tracedb.go
+++ b/server/tracedb/tracedb.go
@@ -17,10 +17,14 @@ type TraceDB interface {
 	ShouldRetry() bool
 	GetTraceID() trace.TraceID
 	GetTraceByID(ctx context.Context, traceID string) (model.Trace, error)
-	TestConnection(ctx context.Context) model.ConnectionResult
 	Close() error
 }
 
+type TestableTraceDB interface {
+	TraceDB
+	TestConnection(ctx context.Context) model.ConnectionResult
+}
+
 type noopTraceDB struct{}
 
 func (db *noopTraceDB) GetTraceByID(ctx context.Context, traceID string) (t model.Trace, err error) {