class TestRedashClient(AppTest):

  def setUp(self):
    api_key = "test_key"
    self.redash = RedashClient(api_key)

    mock_requests_post_patcher = mock.patch(
        "redash_client.client.requests.post")
    self.mock_requests_post = mock_requests_post_patcher.start()
    self.addCleanup(mock_requests_post_patcher.stop)

    mock_requests_get_patcher = mock.patch(
        "redash_client.client.requests.get")
    self.mock_requests_get = mock_requests_get_patcher.start()
    self.addCleanup(mock_requests_get_patcher.stop)

    mock_requests_delete_patcher = mock.patch(
        "redash_client.client.requests.delete")
    self.mock_requests_delete = mock_requests_delete_patcher.start()
    self.addCleanup(mock_requests_delete_patcher.stop)

  def test_request_exception_thrown(self):
    ERROR_STRING = "FAIL"

    def server_call_raising_exception(url, data):
      raise requests.RequestException(ERROR_STRING)

    self.mock_requests_post.side_effect = server_call_raising_exception

    url = "www.test.com"
    self.assertRaisesRegexp(
        self.redash.RedashClientException,
        "Unable to communicate with redash: {0}".format(ERROR_STRING),
        lambda: self.redash._make_request(None, url, args={}))

  def test_failed_request_throws(self):
    STATUS = 404
    ERROR_STRING = "FAIL"
    self.mock_requests_post.return_value = self.get_mock_response(
        STATUS, ERROR_STRING)

    url = "www.test.com"
    self.assertRaisesRegexp(
        self.redash.RedashClientException,
        "Error status returned: {0} {1}".format(STATUS, ERROR_STRING),
        lambda: self.redash._make_request(None, url, args={}))

  def test_failed_to_load_content_json(self):
    BAD_JSON = "boop beep _ epic json fail"
    JSON_ERROR = "No JSON object could be decoded"
    self.mock_requests_post.return_value = self.get_mock_response(
        content=BAD_JSON)

    url = "www.test.com"
    self.assertRaisesRegexp(
        self.redash.RedashClientException,
        "Unable to parse JSON response: {0}".format(JSON_ERROR),
        lambda: self.redash._make_request(None, url, args={}))

  def test_create_new_query_returns_expected_ids(self):
    EXPECTED_QUERY_ID = "query_id123"
    EXPECTED_VIZ_ID = "viz_id123"
    QUERY_ID_RESPONSE = {
        "id": EXPECTED_QUERY_ID
    }

    VISUALIZATION_LIST_RESPONSE = {
        "visualizations": [{
            "id": EXPECTED_VIZ_ID
        }]
    }

    self.mock_requests_post.return_value = self.get_mock_response(
        content=json.dumps(QUERY_ID_RESPONSE))
    self.mock_requests_get.return_value = self.get_mock_response(
        content=json.dumps(VISUALIZATION_LIST_RESPONSE))

    query_id, table_id = self.redash.create_new_query(
        "Dash Name",
        "SELECT * FROM test", 5)

    self.assertEqual(query_id, EXPECTED_QUERY_ID)
    self.assertEqual(table_id, EXPECTED_VIZ_ID)
    self.assertEqual(self.mock_requests_get.call_count, 1)
    self.assertEqual(self.mock_requests_post.call_count, 2)

  def test_create_new_query_returns_none(self):
    QUERY_FAULTY_RESPONSE = {
        "some_bad_response": "boop"
    }

    self.mock_requests_post.return_value = self.get_mock_response(
        content=json.dumps(QUERY_FAULTY_RESPONSE))

    query_id, table_id = self.redash.create_new_query(
        "Dash Name",
        "SELECT * FROM test", 5)

    self.assertEqual(query_id, None)
    self.assertEqual(table_id, None)
    self.assertEqual(self.mock_requests_post.call_count, 1)
    self.assertEqual(self.mock_requests_get.call_count, 0)

  def test_immediate_query_results_are_correct(self):
    EXPECTED_ROWS = [{
        "col1": 123,
        "col2": 456,
    }, {
        "col1": 789,
        "col2": 123,
    }]

    QUERY_RESULTS_RESPONSE = {
        "query_result": {
            "data": {
                "rows": EXPECTED_ROWS
            }
        }
    }

    self.mock_requests_post.return_value = self.get_mock_response(
        content=json.dumps(QUERY_RESULTS_RESPONSE))

    rows = self.redash.get_query_results("SELECT * FROM test", 5)

    self.assertItemsEqual(rows, EXPECTED_ROWS)
    self.assertEqual(self.mock_requests_post.call_count, 1)

  def test_late_response_query_results_are_correct(self):
    EXPECTED_ROWS = [{
        "col1": 123,
        "col2": 456,
    }, {
        "col1": 789,
        "col2": 123,
    }]

    QUERY_RESULTS_RESPONSE = {
        "query_result": {
            "data": {
                "rows": EXPECTED_ROWS
            }
        }
    }
    QUERY_RESULTS_NOT_READY_RESPONSE = {
        "job": {}
    }

    self.server_calls = 0

    def simulate_server_calls(url, data):
      response = QUERY_RESULTS_NOT_READY_RESPONSE
      if self.server_calls >= 2:
        response = QUERY_RESULTS_RESPONSE

      self.server_calls += 1
      return self.get_mock_response(content=json.dumps(response))

    self.mock_requests_post.side_effect = simulate_server_calls

    rows = self.redash.get_query_results("SELECT * FROM test", 5)

    self.assertEqual(rows, EXPECTED_ROWS)
    self.assertEqual(self.mock_requests_post.call_count, 3)

  def test_query_results_not_available(self):
    QUERY_RESULTS_NOT_READY_RESPONSE = {
        "job": {}
    }

    self.mock_requests_post.return_value = self.get_mock_response(
        content=json.dumps(QUERY_RESULTS_NOT_READY_RESPONSE))

    rows = self.redash.get_query_results("SELECT * FROM test", 5)

    self.assertEqual(rows, [])
    self.assertEqual(self.mock_requests_post.call_count, 5)

  def test_new_visualization_throws_for_missing_chart_data(self):
    EXPECTED_QUERY_ID = "query_id123"

    self.assertRaises(ValueError,
                      lambda: self.redash.create_new_visualization(
                          EXPECTED_QUERY_ID, VizType.CHART))

  def test_new_visualization_throws_for_missing_cohort_data(self):
    EXPECTED_QUERY_ID = "query_id123"

    self.assertRaises(ValueError,
                      lambda: self.redash.create_new_visualization(
                          EXPECTED_QUERY_ID, VizType.COHORT))

  def test_new_visualization_throws_for_unexpected_visualization_type(self):
    EXPECTED_QUERY_ID = "query_id123"

    self.assertRaises(ValueError,
                      lambda: self.redash.create_new_visualization(
                          EXPECTED_QUERY_ID, "boop"))

  def test_new_viz_returns_expected_query_id(self):
    EXPECTED_QUERY_ID = "query_id123"
    QUERY_ID_RESPONSE = {
        "id": EXPECTED_QUERY_ID
    }
    TIME_INTERVAL = "weekly"

    self.mock_requests_post.return_value = self.get_mock_response(
        content=json.dumps(QUERY_ID_RESPONSE))

    query_id = self.redash.create_new_visualization(
        EXPECTED_QUERY_ID, VizType.COHORT, time_interval=TIME_INTERVAL)

    self.assertEqual(query_id, EXPECTED_QUERY_ID)
    self.assertEqual(self.mock_requests_post.call_count, 1)

  def test_format_cohort_options_correctly(self):
    TIME_INTERVAL = "weekly"
    COHORT_OPTIONS = {
        "timeInterval": TIME_INTERVAL
    }

    options = self.redash.make_visualization_options(
        viz_type=VizType.COHORT, time_interval=TIME_INTERVAL)
    self.assertItemsEqual(options, COHORT_OPTIONS)

  def test_format_chart_options_correctly(self):
    COLUMN_MAPPING = {"date": "x", "event_rate": "y", "type": "series"}
    CHART_OPTIONS = {
        "globalSeriesType": ChartType.LINE,
        "sortX": True,
        "legend": {"enabled": True},
        "yAxis": [{"type": "linear"}, {"type": "linear", "opposite": True}],
        "series": {"stacking": None},
        "xAxis": {"type": "datetime", "labels": {"enabled": True}},
        "seriesOptions": {},
        "columnMapping": COLUMN_MAPPING,
        "bottomMargin": 50
    }

    options = self.redash.make_visualization_options(
        ChartType.LINE, VizType.CHART, COLUMN_MAPPING)
    self.assertItemsEqual(options, CHART_OPTIONS)

  def test_make_correct_slug(self):
    DASH_NAME = "Activity Stream A/B Testing: Beep Meep"
    EXPECTED_SLUG = "activity-stream-a-b-testing-beep-meep"

    produced_slug = self.redash.get_slug(DASH_NAME)
    self.assertEqual(produced_slug, EXPECTED_SLUG)

  def test_new_dashboard_exists(self):
    DASH_NAME = "Activity Stream A/B Testing: Beep Meep"
    EXPECTED_QUERY_ID = "query_id123"
    QUERY_ID_RESPONSE = {
        "id": EXPECTED_QUERY_ID
    }

    self.mock_requests_get.return_value = self.get_mock_response(
        content=json.dumps(QUERY_ID_RESPONSE))

    query_id = self.redash.create_new_dashboard(DASH_NAME)

    self.assertEqual(query_id, EXPECTED_QUERY_ID)
    self.assertEqual(self.mock_requests_get.call_count, 1)
    self.assertEqual(self.mock_requests_post.call_count, 0)

  def test_new_dashboard_doesnt_exist(self):
    DASH_NAME = "Activity Stream A/B Testing: Beep Meep"
    EXPECTED_QUERY_ID = "query_id123"
    QUERY_ID_RESPONSE = {
        "id": EXPECTED_QUERY_ID
    }

    self.mock_requests_get.return_value = self.get_mock_response(status=404)
    self.mock_requests_post.return_value = self.get_mock_response(
        content=json.dumps(QUERY_ID_RESPONSE))

    query_id = self.redash.create_new_dashboard(DASH_NAME)

    self.assertEqual(query_id, EXPECTED_QUERY_ID)
    self.assertEqual(self.mock_requests_get.call_count, 1)
    self.assertEqual(self.mock_requests_post.call_count, 1)

  def test_publish_dashboard_success(self):
    self.mock_requests_post.return_value = self.get_mock_response()

    self.redash.publish_dashboard(dash_id=1234)

    self.assertEqual(self.mock_requests_post.call_count, 1)
    self.assertEqual(self.mock_requests_get.call_count, 0)

  def test_remove_visualization_success(self):
    self.mock_requests_delete.return_value = self.get_mock_response()

    self.redash.remove_visualization(viz_id=1234)

    self.assertEqual(self.mock_requests_post.call_count, 0)
    self.assertEqual(self.mock_requests_get.call_count, 0)
    self.assertEqual(self.mock_requests_delete.call_count, 1)

  def test_delete_query_success(self):
    self.mock_requests_delete.return_value = self.get_mock_response()

    self.redash.delete_query(query_id=1234)

    self.assertEqual(self.mock_requests_post.call_count, 0)
    self.assertEqual(self.mock_requests_get.call_count, 0)
    self.assertEqual(self.mock_requests_delete.call_count, 1)

  def test_add_visualization_to_dashboard_success(self):
    self.mock_requests_post.return_value = self.get_mock_response()

    self.redash.add_visualization_to_dashboard(
        dash_id=1234, viz_id=5678, viz_width=VizWidth.WIDE)

    self.assertEqual(self.mock_requests_post.call_count, 1)
    self.assertEqual(self.mock_requests_get.call_count, 0)
    self.assertEqual(self.mock_requests_delete.call_count, 0)

  def test_add_visualization_to_dashboard_throws(self):
    self.assertRaises(ValueError,
                      lambda: self.redash.add_visualization_to_dashboard(
                          dash_id=1234, viz_id=5678, viz_width="meep"))

  def test_update_query_schedule_success(self):
    self.mock_requests_post.return_value = self.get_mock_response()

    self.redash.update_query_schedule(query_id=1234, schedule=86400)

    self.assertEqual(self.mock_requests_post.call_count, 1)
    self.assertEqual(self.mock_requests_get.call_count, 0)
    self.assertEqual(self.mock_requests_delete.call_count, 0)

  def test_update_query_string_success(self):
    self.mock_requests_post.return_value = self.get_mock_response()

    self.redash.update_query(
        query_id=1234,
        name="Test",
        sql_query="SELECT * FROM table",
        data_source_id=5,
        description="",
    )

    # One call to update query, one call to refresh it
    self.assertEqual(self.mock_requests_post.call_count, 2)
    self.assertEqual(self.mock_requests_get.call_count, 0)
    self.assertEqual(self.mock_requests_delete.call_count, 0)

  def test_fork_query_returns_correct_attributes(self):
    FORKED_QUERY = {
        "id": 5,
        "query": "sql query text",
        "data_source_id": 5
    }

    self.mock_requests_post.return_value = self.get_mock_response(
        content=json.dumps(FORKED_QUERY))

    fork = self.redash.fork_query(5)

    self.assertEqual(len(fork), 3)
    self.assertTrue("id" in fork)
    self.assertTrue("query" in fork)
    self.assertTrue("data_source_id" in fork)
    self.assertEqual(self.mock_requests_post.call_count, 1)

  def test_search_queries_returns_correct_attributes(self):
    self.get_calls = 0
    QUERIES_IN_SEARCH = [{
        "id": 5,
        "description": "SomeQuery",
        "name": "Query Title",
        "data_source_id": 5
    }]
    VISUALIZATIONS_FOR_QUERY = {
        "visualizations": [
            {"options": {}},
            {"options": {}}
        ]
    }

    def get_server(url):
      response = self.get_mock_response()
      if self.get_calls == 0:
        response = self.get_mock_response(
            content=json.dumps(QUERIES_IN_SEARCH))
      else:
        response = self.get_mock_response(
            content=json.dumps(VISUALIZATIONS_FOR_QUERY))

      self.get_calls += 1
      return response

    self.mock_requests_get.side_effect = get_server

    templates = self.redash.search_queries("Keyword")

    self.assertEqual(len(templates), 1)
    self.assertTrue("id" in templates[0])
    self.assertTrue("description" in templates[0])
    self.assertTrue("name" in templates[0])
    self.assertTrue("data_source_id" in templates[0])
    self.assertEqual(self.mock_requests_get.call_count, 2)

  def test_get_widget_from_dash_returns_correctly_flattened_widgets(self):
    DASH_NAME = "Activity Stream A/B Testing: Beep Meep"
    EXPECTED_QUERY_ID = "query_id123"
    EXPECTED_QUERY_ID2 = "query_id456"
    EXPECTED_QUERY_ID3 = "query_id789"
    FLAT_WIDGETS = [{
        "visualization": {
            "query": {
                "id": EXPECTED_QUERY_ID
            }
        }
    }, {
        "visualization": {
            "query": {
                "id": EXPECTED_QUERY_ID2
            }
        }
    }, {
        "visualization": {
            "query": {
                "id": EXPECTED_QUERY_ID3
            }
        }
    }]

    WIDGETS_RESPONSE = {
        "widgets": [[{
            "visualization": {
                "query": {
                    "id": EXPECTED_QUERY_ID
                }
            }}],
            [{"visualization": {
                "query": {
                    "id": EXPECTED_QUERY_ID2
                }
            }},
            {"visualization": {
                "query": {
                    "id": EXPECTED_QUERY_ID3
                }
            }}
        ]]
    }

    self.mock_requests_get.return_value = self.get_mock_response(
        content=json.dumps(WIDGETS_RESPONSE))

    widget_list = self.redash.get_widget_from_dash(DASH_NAME)

    self.assertEqual(widget_list, FLAT_WIDGETS)
    self.assertEqual(self.mock_requests_get.call_count, 1)
class TestRedashClient(AppTest):
    def setUp(self):
        # Maintain python2 compatibility
        if not hasattr(self, 'assertCountEqual'):  # pragma: no cover
            self.assertCountEqual = self.assertItemsEqual
        if not hasattr(self, 'assertRaisesRegex'):  # pragma: no cover
            self.assertRaisesRegex = self.assertRaisesRegexp

        api_key = "test_key"
        self.redash = RedashClient(api_key)

        mock_requests_post_patcher = mock.patch(
            "redash_client.client.requests.post")
        self.mock_requests_post = mock_requests_post_patcher.start()
        self.addCleanup(mock_requests_post_patcher.stop)

        mock_requests_get_patcher = mock.patch(
            "redash_client.client.requests.get")
        self.mock_requests_get = mock_requests_get_patcher.start()
        self.addCleanup(mock_requests_get_patcher.stop)

        mock_requests_delete_patcher = mock.patch(
            "redash_client.client.requests.delete")
        self.mock_requests_delete = mock_requests_delete_patcher.start()
        self.addCleanup(mock_requests_delete_patcher.stop)

    def test_request_exception_thrown(self):
        ERROR_STRING = "FAIL"

        def server_call_raising_exception(url, data):
            raise requests.RequestException(ERROR_STRING)

        self.mock_requests_post.side_effect = server_call_raising_exception

        url = "www.test.com"
        self.assertRaisesRegex(
            self.redash.RedashClientException,
            "Unable to communicate with redash: {0}".format(ERROR_STRING),
            lambda: self.redash._make_request(None, url, req_args={}))

    def test_failed_request_throws(self):
        STATUS = 404
        ERROR_STRING = "FAIL"
        self.mock_requests_post.return_value = self.get_mock_response(
            STATUS, ERROR_STRING)

        url = "www.test.com"
        self.assertRaisesRegex(
            self.redash.RedashClientException,
            "Error status returned: {0} {1}".format(STATUS, ERROR_STRING),
            lambda: self.redash._make_request(None, url, req_args={}))

    def test_failed_to_load_content_json(self):
        BAD_JSON = "boop beep _ epic json fail"
        JSON_ERROR = "No JSON object could be decoded"
        post_response = self.get_mock_response(content=BAD_JSON)
        post_response.json.side_effect = ValueError(JSON_ERROR)
        self.mock_requests_post.return_value = post_response

        url = "www.test.com"
        self.assertRaisesRegex(
            self.redash.RedashClientException,
            "Unable to parse JSON response: {0}".format(JSON_ERROR),
            lambda: self.redash._make_request(None, url, req_args={}))

    def test_get_public_url_returns_expected_url(self):
        DASH_ID = 6
        EXPECTED_PUBLIC_URL = {"public_url": "www.example.com/expected"}
        post_response = self.get_mock_response(
            content=json.dumps(EXPECTED_PUBLIC_URL))
        post_response.json.return_value = EXPECTED_PUBLIC_URL
        self.mock_requests_post.return_value = post_response

        public_url = self.redash.get_public_url(DASH_ID)
        self.assertEqual(public_url, EXPECTED_PUBLIC_URL["public_url"])

    def test_get_visualization_public_url_has_correct_url(self):
        WIDGET_ID = 123
        QUERY_ID = 456
        URL_PARAM = "api_key={api_key}".format(api_key=self.redash._api_key)

        EXPECTED_PUBLIC_URL = ("https://sql.telemetry.mozilla.org/embed/"
                               "query/{query_id}/visualization/{viz_id}"
                               "?{url_param}").format(query_id=QUERY_ID,
                                                      viz_id=WIDGET_ID,
                                                      url_param=URL_PARAM)

        public_url = self.redash.get_visualization_public_url(
            QUERY_ID, WIDGET_ID)
        self.assertEqual(public_url, EXPECTED_PUBLIC_URL)

    def test_create_new_query_returns_expected_ids(self):
        EXPECTED_QUERY_ID = "query_id123"
        EXPECTED_VIZ_ID = "viz_id123"
        QUERY_ID_RESPONSE = {"id": EXPECTED_QUERY_ID}

        VISUALIZATION_LIST_RESPONSE = {
            "visualizations": [{
                "id": EXPECTED_VIZ_ID
            }]
        }

        query_id_response_json = json.dumps(QUERY_ID_RESPONSE)
        post_response = self.get_mock_response(content=query_id_response_json)
        post_response.json.return_value = QUERY_ID_RESPONSE
        self.mock_requests_post.return_value = post_response

        viz_list_response_json = json.dumps(VISUALIZATION_LIST_RESPONSE)
        get_response = self.get_mock_response(content=viz_list_response_json)
        get_response.json.return_value = VISUALIZATION_LIST_RESPONSE
        self.mock_requests_get.return_value = get_response

        query_id, table_id = self.redash.create_new_query(
            "Dash Name", "SELECT * FROM test", 5)

        self.assertEqual(query_id, EXPECTED_QUERY_ID)
        self.assertEqual(table_id, EXPECTED_VIZ_ID)
        self.assertEqual(self.mock_requests_get.call_count, 1)
        self.assertEqual(self.mock_requests_post.call_count, 2)

    def test_create_new_query_returns_none(self):
        QUERY_FAULTY_RESPONSE = {"some_bad_response": "boop"}
        post_response = self.get_mock_response(
            content=json.dumps(QUERY_FAULTY_RESPONSE))
        post_response.json.return_value = QUERY_FAULTY_RESPONSE
        self.mock_requests_post.return_value = post_response

        query_id, table_id = self.redash.create_new_query(
            "Dash Name", "SELECT * FROM test", 5)

        self.assertEqual(query_id, None)
        self.assertEqual(table_id, None)
        self.assertEqual(self.mock_requests_post.call_count, 1)
        self.assertEqual(self.mock_requests_get.call_count, 0)

    def test_immediate_query_results_are_correct(self):
        EXPECTED_ROWS = [{
            "col1": 123,
            "col2": 456,
        }, {
            "col1": 789,
            "col2": 123,
        }]

        QUERY_RESULTS_RESPONSE = {
            "query_result": {
                "data": {
                    "rows": EXPECTED_ROWS
                }
            }
        }

        post_response = self.get_mock_response(
            content=json.dumps(QUERY_RESULTS_RESPONSE))
        post_response.json.return_value = QUERY_RESULTS_RESPONSE
        self.mock_requests_post.return_value = post_response

        rows = self.redash.get_query_results("SELECT * FROM test", 5)

        self.assertCountEqual(rows, EXPECTED_ROWS)
        self.assertEqual(self.mock_requests_post.call_count, 1)

    def test_late_response_query_results_are_correct(self):
        EXPECTED_ROWS = [{
            "col1": 123,
            "col2": 456,
        }, {
            "col1": 789,
            "col2": 123,
        }]

        QUERY_RESULTS_RESPONSE = {
            "query_result": {
                "data": {
                    "rows": EXPECTED_ROWS
                }
            }
        }
        QUERY_RESULTS_NOT_READY_RESPONSE = {"job": {"status": 1, "id": "123"}}

        QUERY_RESULTS_READY_RESPONSE = {
            "job": {
                "status": 3,
                "id": "123",
                "query_result_id": 456
            }
        }

        # We should have one POST request and two GET requests
        post_response = self.get_mock_response(
            content=json.dumps(QUERY_RESULTS_NOT_READY_RESPONSE))
        post_response.json.return_value = QUERY_RESULTS_NOT_READY_RESPONSE
        self.mock_requests_post.return_value = post_response

        self.get_calls = 0

        def simulate_get_calls(url):
            if self.get_calls == 0:
                self.assertTrue("jobs" in url)
                self.assertTrue("123" in url)
                response = QUERY_RESULTS_READY_RESPONSE
                self.get_calls += 1
            else:
                self.assertTrue("query_results" in url)
                self.assertTrue("456" in url)
                response = QUERY_RESULTS_RESPONSE

            get_response = self.get_mock_response(content=json.dumps(response))
            get_response.json.return_value = response
            return get_response

        self.mock_requests_get.side_effect = simulate_get_calls

        rows = self.redash.get_query_results("SELECT * FROM test", 5)

        self.assertEqual(rows, EXPECTED_ROWS)
        self.assertEqual(self.mock_requests_post.call_count, 1)
        self.assertEqual(self.mock_requests_get.call_count, 2)

    def test_query_results_not_available(self):
        QUERY_RESULTS_NOT_READY_RESPONSE = {"job": {"status": 1, "id": "123"}}

        self.redash._retry_delay = .000000001

        post_response = self.get_mock_response(
            content=json.dumps(QUERY_RESULTS_NOT_READY_RESPONSE))
        post_response.json.return_value = QUERY_RESULTS_NOT_READY_RESPONSE
        self.mock_requests_post.return_value = post_response

        get_response = self.get_mock_response(
            content=json.dumps(QUERY_RESULTS_NOT_READY_RESPONSE))
        get_response.json.return_value = QUERY_RESULTS_NOT_READY_RESPONSE
        self.mock_requests_get.return_value = get_response

        rows = self.redash.get_query_results("SELECT * FROM test", 5)

        self.assertEqual(rows, [])
        self.assertEqual(self.mock_requests_post.call_count, 1)
        self.assertEqual(self.mock_requests_get.call_count, 5)

    def test_new_visualization_throws_for_missing_chart_data(self):
        EXPECTED_QUERY_ID = "query_id123"

        self.assertRaises(
            ValueError, lambda: self.redash.create_new_visualization(
                EXPECTED_QUERY_ID, VizType.CHART))

    def test_new_visualization_throws_for_missing_cohort_data(self):
        EXPECTED_QUERY_ID = "query_id123"

        self.assertRaises(
            ValueError, lambda: self.redash.create_new_visualization(
                EXPECTED_QUERY_ID, VizType.COHORT))

    def test_new_visualization_throws_for_unexpected_visualization_type(self):
        EXPECTED_QUERY_ID = "query_id123"

        self.assertRaises(
            ValueError, lambda: self.redash.create_new_visualization(
                EXPECTED_QUERY_ID, "boop"))

    def test_new_viz_returns_expected_query_id(self):
        EXPECTED_QUERY_ID = "query_id123"
        QUERY_ID_RESPONSE = {"id": EXPECTED_QUERY_ID}
        TIME_INTERVAL = "weekly"

        post_response = self.get_mock_response(
            content=json.dumps(QUERY_ID_RESPONSE))
        post_response.json.return_value = QUERY_ID_RESPONSE
        self.mock_requests_post.return_value = post_response

        query_id = self.redash.create_new_visualization(
            EXPECTED_QUERY_ID, VizType.COHORT, time_interval=TIME_INTERVAL)

        self.assertEqual(query_id, EXPECTED_QUERY_ID)
        self.assertEqual(self.mock_requests_post.call_count, 1)

    def test_format_cohort_options_correctly(self):
        TIME_INTERVAL = "weekly"
        COHORT_OPTIONS = {"timeInterval": TIME_INTERVAL}

        options = self.redash.make_visualization_options(
            viz_type=VizType.COHORT, time_interval=TIME_INTERVAL)
        self.assertCountEqual(options, COHORT_OPTIONS)

    def test_format_chart_options_correctly(self):
        COLUMN_MAPPING = {"date": "x", "event_rate": "y", "type": "series"}
        CHART_OPTIONS = {
            "globalSeriesType": ChartType.LINE,
            "sortX": True,
            "legend": {
                "enabled": True
            },
            "yAxis": [{
                "type": "linear"
            }, {
                "type": "linear",
                "opposite": True
            }],
            "series": {
                "stacking": None
            },
            "xAxis": {
                "type": "datetime",
                "labels": {
                    "enabled": True
                }
            },
            "seriesOptions": {},
            "columnMapping": COLUMN_MAPPING,
            "bottomMargin": 50
        }

        options = self.redash.make_visualization_options(
            ChartType.LINE, VizType.CHART, COLUMN_MAPPING)
        self.assertCountEqual(options, CHART_OPTIONS)

    def test_make_correct_slug(self):
        DASH_NAME = "Activity Stream A/B Testing: Beep Meep"
        EXPECTED_SLUG = "activity-stream-a-b-testing-beep-meep"

        produced_slug = self.redash.get_slug(DASH_NAME)
        self.assertEqual(produced_slug, EXPECTED_SLUG)

    def test_new_dashboard_exists(self):
        DASH_NAME = "Activity Stream A/B Testing: Beep Meep"
        EXPECTED_QUERY_ID = "query_id123"
        EXPECTED_SLUG = "some_slug_it_made"
        QUERY_ID_RESPONSE = {"id": EXPECTED_QUERY_ID, "slug": EXPECTED_SLUG}

        get_response = self.get_mock_response(
            content=json.dumps(QUERY_ID_RESPONSE))
        get_response.json.return_value = QUERY_ID_RESPONSE
        self.mock_requests_get.return_value = get_response

        dash_info = self.redash.create_new_dashboard(DASH_NAME)

        self.assertEqual(dash_info["dashboard_id"], EXPECTED_QUERY_ID)
        self.assertEqual(dash_info["dashboard_slug"], EXPECTED_SLUG)
        self.assertEqual(
            dash_info["slug_url"], self.redash.BASE_URL +
            "dashboard/{slug}".format(slug=EXPECTED_SLUG))
        self.assertEqual(self.mock_requests_get.call_count, 1)
        self.assertEqual(self.mock_requests_post.call_count, 0)

    def test_new_dashboard_doesnt_exist(self):
        DASH_NAME = "Activity Stream A/B Testing: Beep Meep"
        EXPECTED_QUERY_ID = "query_id123"
        EXPECTED_SLUG = "some_slug_it_made"
        QUERY_ID_RESPONSE = {"id": EXPECTED_QUERY_ID, "slug": EXPECTED_SLUG}

        self.mock_requests_get.return_value = self.get_mock_response(
            status=404)
        post_response = self.get_mock_response(
            content=json.dumps(QUERY_ID_RESPONSE))
        post_response.json.return_value = QUERY_ID_RESPONSE
        self.mock_requests_post.return_value = post_response

        dash_info = self.redash.create_new_dashboard(DASH_NAME)

        self.assertEqual(dash_info["dashboard_id"], EXPECTED_QUERY_ID)
        self.assertEqual(dash_info["dashboard_slug"], EXPECTED_SLUG)
        self.assertEqual(
            dash_info["slug_url"], self.redash.BASE_URL +
            "dashboard/{slug}".format(slug=EXPECTED_SLUG))
        self.assertEqual(self.mock_requests_get.call_count, 1)
        self.assertEqual(self.mock_requests_post.call_count, 1)

    def test_publish_dashboard_success(self):
        self.mock_requests_post.return_value = self.get_mock_response()

        self.redash.publish_dashboard(dash_id=1234)

        self.assertEqual(self.mock_requests_post.call_count, 1)
        self.assertEqual(self.mock_requests_get.call_count, 0)

    def test_remove_visualization_success(self):
        self.mock_requests_delete.return_value = self.get_mock_response()

        self.redash.remove_visualization(viz_id=1234)

        self.assertEqual(self.mock_requests_post.call_count, 0)
        self.assertEqual(self.mock_requests_get.call_count, 0)
        self.assertEqual(self.mock_requests_delete.call_count, 1)

    def test_delete_query_success(self):
        self.mock_requests_delete.return_value = self.get_mock_response()

        self.redash.delete_query(query_id=1234)

        self.assertEqual(self.mock_requests_post.call_count, 0)
        self.assertEqual(self.mock_requests_get.call_count, 0)
        self.assertEqual(self.mock_requests_delete.call_count, 1)

    def test_add_visualization_to_dashboard_success(self):
        self.mock_requests_post.return_value = self.get_mock_response()

        self.redash.add_visualization_to_dashboard(dash_id=1234,
                                                   viz_id=5678,
                                                   viz_width=VizWidth.WIDE)

        self.assertEqual(self.mock_requests_post.call_count, 1)
        self.assertEqual(self.mock_requests_get.call_count, 0)
        self.assertEqual(self.mock_requests_delete.call_count, 0)

    def test_add_visualization_to_dashboard_throws(self):
        self.assertRaises(
            ValueError, lambda: self.redash.add_visualization_to_dashboard(
                dash_id=1234, viz_id=5678, viz_width="meep"))

    def test_update_query_schedule_success(self):
        self.mock_requests_post.return_value = self.get_mock_response()

        self.redash.update_query_schedule(query_id=1234, schedule=86400)

        self.assertEqual(self.mock_requests_post.call_count, 1)
        self.assertEqual(self.mock_requests_get.call_count, 0)
        self.assertEqual(self.mock_requests_delete.call_count, 0)

    def test_update_query_string_success(self):
        self.mock_requests_post.return_value = self.get_mock_response()

        self.redash.update_query(query_id=1234,
                                 name="Test",
                                 sql_query="SELECT * FROM table",
                                 data_source_id=5,
                                 description="",
                                 options={"some_options": "an_option"})

        # One call to update query, one call to refresh it
        self.assertEqual(self.mock_requests_post.call_count, 2)
        self.assertEqual(self.mock_requests_get.call_count, 0)
        self.assertEqual(self.mock_requests_delete.call_count, 0)

    def test_fork_query_returns_correct_attributes(self):
        FORKED_QUERY = {
            "id": 5,
            "query": "sql query text",
            "data_source_id": 5
        }

        self.mock_requests_post.return_value = self.get_mock_response(
            content=json.dumps(FORKED_QUERY))

        fork = self.redash.fork_query(5)

        self.assertEqual(len(fork), 3)
        self.assertTrue("id" in fork)
        self.assertTrue("query" in fork)
        self.assertTrue("data_source_id" in fork)
        self.assertEqual(self.mock_requests_post.call_count, 1)

    def test_search_queries_returns_correct_attributes(self):
        self.get_calls = 0
        QUERIES_IN_SEARCH = {
            "results": [{
                "id": 5,
                "description": "SomeQuery",
                "name": "Query Title",
                "data_source_id": 5
            }]
        }
        VISUALIZATIONS_FOR_QUERY = {
            "visualizations": [{
                "options": {}
            }, {
                "options": {}
            }]
        }

        def get_server(url):
            response = self.get_mock_response()
            response.json.return_value = {}
            if self.get_calls == 0:
                response = self.get_mock_response(
                    content=json.dumps(QUERIES_IN_SEARCH))
                response.json.return_value = QUERIES_IN_SEARCH
            else:
                response = self.get_mock_response(
                    content=json.dumps(VISUALIZATIONS_FOR_QUERY))
                response.json.return_value = VISUALIZATIONS_FOR_QUERY

            self.get_calls += 1
            return response

        self.mock_requests_get.side_effect = get_server

        templates = self.redash.search_queries("Keyword")

        self.assertEqual(len(templates), 1)
        self.assertTrue("id" in templates[0])
        self.assertTrue("description" in templates[0])
        self.assertTrue("name" in templates[0])
        self.assertTrue("data_source_id" in templates[0])
        self.assertEqual(self.mock_requests_get.call_count, 2)

    def test_get_widget_from_dash_returns_correctly_flattened_widgets(self):
        DASH_NAME = "Activity Stream A/B Testing: Beep Meep"
        EXPECTED_QUERY_ID = "query_id123"
        EXPECTED_QUERY_ID2 = "query_id456"
        EXPECTED_QUERY_ID3 = "query_id789"
        FLAT_WIDGETS = [{
            "visualization": {
                "query": {
                    "id": EXPECTED_QUERY_ID
                }
            }
        }, {
            "visualization": {
                "query": {
                    "id": EXPECTED_QUERY_ID2
                }
            }
        }, {
            "visualization": {
                "query": {
                    "id": EXPECTED_QUERY_ID3
                }
            }
        }]

        WIDGETS_RESPONSE = {
            "widgets": [{
                "visualization": {
                    "query": {
                        "id": EXPECTED_QUERY_ID
                    }
                }
            }, {
                "visualization": {
                    "query": {
                        "id": EXPECTED_QUERY_ID2
                    }
                }
            }, {
                "visualization": {
                    "query": {
                        "id": EXPECTED_QUERY_ID3
                    }
                }
            }]
        }

        get_response = self.get_mock_response(
            content=json.dumps(WIDGETS_RESPONSE))
        get_response.json.return_value = WIDGETS_RESPONSE
        self.mock_requests_get.return_value = get_response

        widget_list = self.redash.get_widget_from_dash(DASH_NAME)

        self.assertEqual(widget_list, FLAT_WIDGETS)
        self.assertEqual(self.mock_requests_get.call_count, 1)

    def test_get_data_sources(self):
        DATA_SOURCES = [{"name": "data_source_1"}, {"name": "data_source_2"}]
        get_response = self.get_mock_response(content=json.dumps(DATA_SOURCES))
        get_response.json.return_value = DATA_SOURCES
        self.mock_requests_get.return_value = get_response

        sources = self.redash.get_data_sources()
        self.assertEqual(sources, DATA_SOURCES)
Beispiel #3
0
class SummaryDashboard(object):
    class ExternalAPIError(Exception):
        pass

    def __init__(self, api_key, dash_name):
        self._dash_name = dash_name

        try:
            self.redash = RedashClient(api_key)

            dash_info = self.redash.create_new_dashboard(self._dash_name)
            self._dash_id = dash_info["dashboard_id"]
            self.slug_url = dash_info["slug_url"]

            self.redash.publish_dashboard(self._dash_id)
            self.public_url = self.redash.get_public_url(self._dash_id)
        except self.redash.RedashClientException as e:
            raise self.ExternalAPIError(
                "Unable to create new dashboard: {error}".format(error=e), e)

    def _create_new_query(self,
                          query_title,
                          query_string,
                          data_source,
                          description=""):
        try:
            query_id, table_id = self.redash.create_new_query(query_title,
                                                              query_string,
                                                              data_source,
                                                              description=None)
            return query_id, table_id
        except self.redash.RedashClientException as e:
            raise self.ExternalAPIError(
                "Unable to create query titled '{title}': {error}".format(
                    title=query_title, error=e))

    def _add_visualization_to_dashboard(self, viz_id, visualization_width):
        try:
            self.redash.add_visualization_to_dashboard(self._dash_id, viz_id,
                                                       visualization_width)
        except self.redash.RedashClientException as e:
            raise self.ExternalAPIError(
                ("Unable to add visualization '{id}' to "
                 "dashboard '{title}': {error}").format(id=viz_id,
                                                        title=self._dash_name,
                                                        error=e))

    def _get_query_results(self, query_string, data_source_id, query_name=""):
        try:
            data = self.redash.get_query_results(query_string, data_source_id)
            return data
        except self.redash.RedashClientException as e:
            raise self.ExternalAPIError(
                "Unable to fetch query results: '{query_name}' "
                " {error}".format(query_name=query_name, error=e))

    def _create_new_visualization(self, query_id, visualization_type,
                                  visualization_name, chart_type,
                                  column_mapping, series_options,
                                  time_interval, stacking, axis_info):
        try:
            viz_id = self.redash.create_new_visualization(
                query_id,
                visualization_type,
                visualization_name,
                chart_type,
                column_mapping,
                series_options,
                time_interval,
                stacking,
                axis_info,
            )
            return viz_id
        except self.redash.RedashClientException as e:
            raise self.ExternalAPIError(
                "Unable to create visualization titled '{title}': {error}".
                format(title=visualization_name, error=e))

    def _get_widgets_from_dash(self, dash_name):
        try:
            return self.redash.get_widget_from_dash(dash_name)
        except self.redash.RedashClientException as e:
            raise self.ExternalAPIError(
                "Unable to access dashboard widgets: {error}".format(error=e),
                e)

    def _update_query(self,
                      query_id,
                      query_title,
                      sql,
                      data_source_id,
                      description="",
                      options=""):
        try:
            self.redash.update_query(
                query_id,
                query_title,
                sql,
                data_source_id,
                description,
                options,
            )
        except self.redash.RedashClientException as e:
            raise self.ExternalAPIError(
                "Unable to update query {title}: {error}".format(
                    title=query_title, error=e))

    def update_refresh_schedule(self, seconds_to_refresh):
        widgets = self._get_widgets_from_dash(self._dash_name)

        for widget in widgets:
            widget_id = widget.get("visualization",
                                   {}).get("query", {}).get("id", None)

            if not widget_id:
                continue

            try:
                self.redash.update_query_schedule(widget_id,
                                                  seconds_to_refresh)
            except self.redash.RedashClientException as e:
                raise self.ExternalAPIError(
                    "Unable to update schedule for widget {widget_id}: {error}"
                    .format(widget_id=widget_id, error=e))

    def get_update_range(self):
        query_data = self.get_query_ids_and_names()

        if len(query_data) < 1:
            return {}

        dates = [
            parse(query_data[graph]["updated_at"]) for graph in query_data
        ]
        update_range = {"min": min(dates), "max": max(dates)}
        return update_range

    def get_query_ids_and_names(self):
        widgets = self._get_widgets_from_dash(self._dash_name)

        data = {}
        for widget in widgets:
            widget_id = widget.get("id", None)

            query_id = widget.get("visualization", {}).get("query",
                                                           {}).get("id", None)

            widget_name = widget.get("visualization",
                                     {}).get("query", {}).get("name", None)

            widget_query = widget.get("visualization",
                                      {}).get("query", {}).get("query", None)

            updated_at = widget.get("visualization",
                                    {}).get("query",
                                            {}).get("updated_at", None)

            if not widget_name:
                continue

            data[widget_name] = {
                "query_id": query_id,
                "widget_id": widget_id,
                "query": widget_query,
                "updated_at": updated_at
            }

        return data

    def remove_graph_from_dashboard(self, widget_id, query_id):
        try:
            if widget_id is not None:
                self.redash.remove_visualization(widget_id)

            if query_id is not None:
                self.redash.delete_query(query_id)
        except self.redash.RedashClientException as e:
            raise self.ExternalAPIError(
                "Unable to remove widget {widget_id} with query ID "
                "{query_id} from dashboard: {error}".format(
                    widget_id=widget_id, query_id=query_id, error=e))

    def remove_all_graphs(self):
        widgets = self.get_query_ids_and_names()

        for widget_name in widgets:
            widget = widgets[widget_name]
            widget_id = widget.get("widget_id", None)
            query_id = widget.get("query_id", None)

            self.remove_graph_from_dashboard(widget_id, query_id)

    def _populate_sql_string_with_variables(self, template_sql, query_params):
        adjusted_string = template_sql.replace("{{{", "{").replace("}}}", "}")
        sql_query = adjusted_string.format(**query_params)
        return sql_query

    def _add_copied_query_to_dashboard(self,
                                       template,
                                       query_title,
                                       query_params,
                                       visualization_width,
                                       visualization_name="Chart"):
        query_string = self._populate_sql_string_with_variables(
            template["query"], query_params)

        query_id, table_id = self._create_new_query(query_title, query_string,
                                                    template["data_source_id"])
        try:
            viz_id = self.redash.make_new_visualization_request(
                query_id,
                template["type"],
                template["options"],
                visualization_name,
            )
            self._add_visualization_to_dashboard(viz_id, visualization_width)

            public_url = self.redash.get_visualization_public_url(
                query_id, viz_id)

            return public_url
        except self.redash.RedashClientException as e:
            raise self.ExternalAPIError(
                "Unable to add copied query {query_id} to "
                "dashboard: {error}".format(query_id=query_id, error=e))

    def _template_copy_results_exist(self, query_title, template_sql,
                                     data_source_id, query_params):
        sql_query = self._populate_sql_string_with_variables(
            template_sql, query_params)
        data = self._get_query_results(sql_query, data_source_id, query_title)

        if data is None or len(data) == 0:
            self._logger.info(
                ("Dashboard: Query '{name}' is still updating and will "
                 "not be displayed.".format(name=query_title)))
            return False

        return True

    def add_query_to_dashboard(self,
                               query_title,
                               query_string,
                               data_source,
                               visualization_width,
                               visualization_type=VizType.CHART,
                               visualization_name="",
                               chart_type=None,
                               column_mapping=None,
                               series_options=None,
                               time_interval=None,
                               stacking=True,
                               axis_info={}):
        query_id, table_id = self._create_new_query(query_title, query_string,
                                                    data_source)
        viz_id = self._create_new_visualization(
            query_id,
            visualization_type,
            visualization_name,
            chart_type,
            column_mapping,
            series_options,
            time_interval,
            stacking,
            axis_info,
        )
        self._add_visualization_to_dashboard(viz_id, visualization_width)