Fix support for custom visualizations

backend/src/apiserver/server/visualization_server.go

@@ -39,7 +39,8 @@ func (s *VisualizationServer) CreateVisualization(ctx context.Context, request *
 // It returns an error if a go_client.Visualization object does not have valid
 // values.
 func (s *VisualizationServer) validateCreateVisualizationRequest(request *go_client.CreateVisualizationRequest) error {
-	if len(request.Visualization.Source) == 0 {
+	// Only validate that a source is provided for non-custom visualizations.
+	if request.Visualization.Type != go_client.Visualization_CUSTOM && len(request.Visualization.Source) == 0 {
 		return util.NewInvalidInputError("A visualization requires a Source to be provided. Received %s", request.Visualization.Source)
 	}
 	// Manually set Arguments to empty JSON if nothing is provided. This is done
@@ -67,7 +68,13 @@ func (s *VisualizationServer) generateVisualizationFromRequest(request *go_clien
 		)
 	}
 	visualizationType := strings.ToLower(go_client.Visualization_Type_name[int32(request.Visualization.Type)])
-	arguments := fmt.Sprintf("--type %s --source %s --arguments '%s'", visualizationType, request.Visualization.Source, request.Visualization.Arguments)
+	arguments := fmt.Sprintf("--type %s --arguments '''%s'''", visualizationType, request.Visualization.Arguments)
+	if !(request.Visualization.Type == go_client.Visualization_CUSTOM && len(request.Visualization.Source) == 0) {
+		// Only add the source argument if a visualization is one of the following:
+		// - Not a custom visualization
+		// - A custom visualization that has a source specified
+		arguments += fmt.Sprintf(" --source %s", request.Visualization.Source)
+	}
 	resp, err := http.PostForm(s.serviceURL, url.Values{"arguments": {arguments}})
 	if err != nil {
 		return nil, util.Wrap(err, "Unable to initialize visualization request.")

backend/src/apiserver/server/visualization_server_test.go

@@ -65,6 +65,24 @@ func TestValidateCreateVisualizationRequest_SourceIsEmpty(t *testing.T) {
 	assert.Contains(t, err.Error(), "A visualization requires a Source to be provided. Received")
 }
 
+func TestValidateCreateVisualizationRequest_SourceIsEmptyAndTypeIsCustom(t *testing.T) {
+	clients, manager, _ := initWithExperiment(t)
+	defer clients.Close()
+	server := &VisualizationServer{
+		resourceManager:    manager,
+		isServiceAvailable: false,
+	}
+	visualization := &go_client.Visualization{
+		Type:      go_client.Visualization_CUSTOM,
+		Arguments: "{}",
+	}
+	request := &go_client.CreateVisualizationRequest{
+		Visualization: visualization,
+	}
+	err := server.validateCreateVisualizationRequest(request)
+	assert.Nil(t, err)
+}
+
 func TestValidateCreateVisualizationRequest_ArgumentsNotValidJSON(t *testing.T) {
 	clients, manager, _ := initWithExperiment(t)
 	defer clients.Close()

backend/src/apiserver/visualization/exporter.py

@@ -118,6 +118,19 @@ def create_cell_from_file(filepath: Text) -> NotebookNode:
 
         return new_code_cell(code)
 
+    @staticmethod
+    def create_cell_from_custom_code(code: list) -> NotebookNode:
+        """Creates a NotebookNode object with provided list as code in node.
+
+        Args:
+            code: list representing lines of code to be run.
+
+        Returns:
+            NotebookNode with specified file as code within node.
+
+        """
+        return new_code_cell("\n".join(code))
+
     def generate_html_from_notebook(self, nb: NotebookNode) -> Text:
         """Converts a provided NotebookNode to HTML.
 

backend/src/apiserver/visualization/server.py

@@ -97,7 +97,7 @@ def is_valid_request_arguments(self, arguments: Namespace):
         """
         if arguments.type is None:
             raise Exception("No type specified.")
-        if arguments.source is None:
+        if arguments.type != "custom" and arguments.source is None:
             raise Exception("No source specified.")
         try:
             json.loads(arguments.arguments)
@@ -126,8 +126,13 @@ def generate_notebook_from_arguments(
         nb = new_notebook()
         nb.cells.append(_exporter.create_cell_from_args(arguments))
         nb.cells.append(new_code_cell('source = "{}"'.format(source)))
-        visualization_file = str(Path.cwd() / "types/{}.py".format(visualization_type))
-        nb.cells.append(_exporter.create_cell_from_file(visualization_file))
+        if visualization_type == "custom":
+            code = arguments.get("code", [])
+            nb.cells.append(_exporter.create_cell_from_custom_code(code))
+        else:
+            visualization_file = str(Path.cwd() / "types/{}.py".format(visualization_type))
+            nb.cells.append(_exporter.create_cell_from_file(visualization_file))
+
         return nb
 
     def get(self):

backend/src/apiserver/visualization/snapshots/snap_test_exporter.py

@@ -122,3 +122,6 @@
 
 
 '''
+
+snapshots['TestExporterMethods::test_create_cell_from_custom_code 1'] = '''x = 2
+print(x)'''

backend/src/apiserver/visualization/test_exporter.py

@@ -61,6 +61,15 @@ def test_create_cell_from_file(self):
         cell = self.exporter.create_cell_from_file("types/test.py")
         self.assertMatchSnapshot(cell.source)
 
+    def test_create_cell_from_custom_code(self):
+        self.maxDiff = None
+        code = [
+            "x = 2",
+            "print(x)"
+        ]
+        cell = self.exporter.create_cell_from_custom_code(code)
+        self.assertMatchSnapshot(cell.source)
+
     def test_generate_html_from_notebook(self):
         self.maxDiff = None
         nb = new_notebook()

backend/src/apiserver/visualization/test_server.py

@@ -70,6 +70,13 @@ def test_create_visualization_fails_when_missing_input_path(self):
             response.body
         )
 
+    def test_create_visualization_passes_when_missing_input_path_and_type_is_custom(self):
+        response = self.fetch(
+            "/",
+            method="POST",
+            body='arguments=--type custom')
+        self.assertEqual(200, response.code)
+
     def test_create_visualization_fails_when_invalid_json_is_provided(self):
         response = self.fetch(
             "/",