From fcedccc436220528174292e9542c3a5b7ccac073 Mon Sep 17 00:00:00 2001
From: Kirin Patel <kirinpatel@gmail.com>
Date: Thu, 15 Aug 2019 16:51:55 -0700
Subject: [PATCH] Changed host name to be constant provided by environment

---
 backend/src/apiserver/client_manager.go       |  2 ++
 backend/src/apiserver/main.go                 |  2 +-
 .../apiserver/server/visualization_server.go  |  4 ++--
 .../server/visualization_server_test.go       | 20 +++++++++++++++----
 4 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/backend/src/apiserver/client_manager.go b/backend/src/apiserver/client_manager.go
index 7004857380c5..5c24014a3278 100644
--- a/backend/src/apiserver/client_manager.go
+++ b/backend/src/apiserver/client_manager.go
@@ -51,6 +51,8 @@ const (
 	podNamespace          = "POD_NAMESPACE"
 	dbName                = "mlpipeline"
 	initConnectionTimeout = "InitConnectionTimeout"
+
+	visualizationServiceHost = "ML_VISUALIZATION_SERVICE_HOST"
 )
 
 // Container for all service clients
diff --git a/backend/src/apiserver/main.go b/backend/src/apiserver/main.go
index 0ddca77b51c6..13e7be124799 100644
--- a/backend/src/apiserver/main.go
+++ b/backend/src/apiserver/main.go
@@ -83,7 +83,7 @@ func startRpcServer(resourceManager *resource.ResourceManager) {
 	api.RegisterRunServiceServer(s, server.NewRunServer(resourceManager))
 	api.RegisterJobServiceServer(s, server.NewJobServer(resourceManager))
 	api.RegisterReportServiceServer(s, server.NewReportServer(resourceManager))
-	api.RegisterVisualizationServiceServer(s, server.NewVisualizationServer(resourceManager, getStringConfig(podNamespace)))
+	api.RegisterVisualizationServiceServer(s, server.NewVisualizationServer(resourceManager, getStringConfig(visualizationServiceHost), getStringConfig(podNamespace)))
 
 	// Register reflection service on gRPC server.
 	reflection.Register(s)
diff --git a/backend/src/apiserver/server/visualization_server.go b/backend/src/apiserver/server/visualization_server.go
index fce341f1e6e1..a74bb3153199 100644
--- a/backend/src/apiserver/server/visualization_server.go
+++ b/backend/src/apiserver/server/visualization_server.go
@@ -80,8 +80,8 @@ func (s *VisualizationServer) generateVisualizationFromRequest(request *go_clien
 	return body, nil
 }
 
-func NewVisualizationServer(resourceManager *resource.ResourceManager, namespace string) *VisualizationServer {
-	serviceURL := fmt.Sprintf("http://ml-pipeline-visualization.%s", namespace)
+func NewVisualizationServer(resourceManager *resource.ResourceManager, serviceName string, namespace string) *VisualizationServer {
+	serviceURL := fmt.Sprintf("http://%s.%s", serviceName, namespace)
 	_, err := http.Get(serviceURL)
 	return &VisualizationServer{
 		resourceManager:    resourceManager,
diff --git a/backend/src/apiserver/server/visualization_server_test.go b/backend/src/apiserver/server/visualization_server_test.go
index 214a6d2186a4..82c39e270128 100644
--- a/backend/src/apiserver/server/visualization_server_test.go
+++ b/backend/src/apiserver/server/visualization_server_test.go
@@ -11,7 +11,10 @@ import (
 func TestValidateCreateVisualizationRequest(t *testing.T) {
 	clients, manager, _ := initWithExperiment(t)
 	defer clients.Close()
-	server := NewVisualizationServer(manager, "kubeflow")
+	server := &VisualizationServer{
+		resourceManager: manager,
+		isServiceAvailable: false,
+	}
 	visualization := &go_client.Visualization{
 		Type:      go_client.Visualization_ROC_CURVE,
 		Source: "gs://ml-pipeline/roc/data.csv",
@@ -27,7 +30,10 @@ func TestValidateCreateVisualizationRequest(t *testing.T) {
 func TestValidateCreateVisualizationRequest_ArgumentsAreEmpty(t *testing.T) {
 	clients, manager, _ := initWithExperiment(t)
 	defer clients.Close()
-	server := NewVisualizationServer(manager, "kubeflow")
+	server := &VisualizationServer{
+		resourceManager: manager,
+		isServiceAvailable: false,
+	}
 	visualization := &go_client.Visualization{
 		Type:      go_client.Visualization_ROC_CURVE,
 		Source: "gs://ml-pipeline/roc/data.csv",
@@ -43,7 +49,10 @@ func TestValidateCreateVisualizationRequest_ArgumentsAreEmpty(t *testing.T) {
 func TestValidateCreateVisualizationRequest_SourceIsEmpty(t *testing.T) {
 	clients, manager, _ := initWithExperiment(t)
 	defer clients.Close()
-	server := NewVisualizationServer(manager, "kubeflow")
+	server := &VisualizationServer{
+		resourceManager: manager,
+		isServiceAvailable: false,
+	}
 	visualization := &go_client.Visualization{
 		Type:      go_client.Visualization_ROC_CURVE,
 		Source: "",
@@ -59,7 +68,10 @@ func TestValidateCreateVisualizationRequest_SourceIsEmpty(t *testing.T) {
 func TestValidateCreateVisualizationRequest_ArgumentsNotValidJSON(t *testing.T) {
 	clients, manager, _ := initWithExperiment(t)
 	defer clients.Close()
-	server := NewVisualizationServer(manager, "kubeflow")
+	server := &VisualizationServer{
+		resourceManager: manager,
+		isServiceAvailable: false,
+	}
 	visualization := &go_client.Visualization{
 		Type:      go_client.Visualization_ROC_CURVE,
 		Source: "gs://ml-pipeline/roc/data.csv",