Read stats source from db

EFForg · Aug 21, 2019 · ed3f481 · ed3f481
1 parent 003fd9a
commit ed3f481
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/api/stats_test.go b/api/stats_test.go
@@ -21,6 +21,15 @@ func TestGetStats(t *testing.T) {
 		MTASTSEnforce: 2,
 	})
 
+	err = api.Database.PutAggregatedScan(checker.AggregatedScan{
+		Time:          time.Now(),
+		Source:        checker.TopDomainsSource,
+		Attempted:     10,
+		WithMXs:       8,
+		MTASTSTesting: 3,
+		MTASTSEnforce: 2,
+	})
+
 	resp, err := http.Get(server.URL + "/api/stats")
 	if err != nil {
 		t.Fatal(err)
@@ -36,8 +45,14 @@ func TestGetStats(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	// Local source returns a percent
 	expectedY := fmt.Sprintf("\"y\": 62.5")
 	if !strings.Contains(string(body), expectedY) {
 		t.Errorf("Expected %s to contain %s", string(body), expectedY)
 	}
+	// Top domains source returns a raw count
+	expectedY = fmt.Sprintf("\"y\": 5")
+	if !strings.Contains(string(body), expectedY) {
+		t.Errorf("Expected %s to contain %s", string(body), expectedY)
+	}
 }
diff --git a/db/sqldb.go b/db/sqldb.go
@@ -122,7 +122,7 @@ func (db *SQLDatabase) PutScan(scan models.Scan) error {
 func (db *SQLDatabase) GetStats(source string) (stats.Series, error) {
 	series := stats.Series{}
 	rows, err := db.conn.Query(
-		`SELECT time, with_mxs, mta_sts_testing, mta_sts_enforce
+		`SELECT time, source, with_mxs, mta_sts_testing, mta_sts_enforce
 		FROM aggregated_scans
 		WHERE source=$1
 		ORDER BY time`, source)
@@ -132,7 +132,7 @@ func (db *SQLDatabase) GetStats(source string) (stats.Series, error) {
 	defer rows.Close()
 	for rows.Next() {
 		var a checker.AggregatedScan
-		if err := rows.Scan(&a.Time, &a.WithMXs, &a.MTASTSTesting, &a.MTASTSEnforce); err != nil {
+		if err := rows.Scan(&a.Time, &a.Source, &a.WithMXs, &a.MTASTSTesting, &a.MTASTSEnforce); err != nil {
 			return series, err
 		}
 		series = append(series, a)