entur: Only include lines operated by AtB

3 files changed, 32 insertions, 2 deletions

diff --git a/entur/entur.go b/entur/entur.go
index d918ca3..d68c764 100644
--- a/entur/entur.go
+++ b/entur/entur.go
@@ -61,9 +61,14 @@ type destinationDisplay struct {
 }
 
 type serviceJourney struct {
+	Operator operator `json:"operator"`
 	JourneyPattern journeyPattern `json:"journeyPattern"`
 }
 
+type operator struct {
+	Id string `json:"id"`
+}
+
 type journeyPattern struct {
 	DirectionType string `json:"directionType"`
 	Line          line   `json:"line"`
@@ -76,7 +81,7 @@ type line struct {
 // Departures returns departures from the given stop ID. Use https://stoppested.entur.org/ to determine stop IDs.
 func (c *Client) Departures(count, stopID int) ([]Departure, error) {
 	// https://api.entur.io/journey-planner/v2/ide/ for query testing
-	query := fmt.Sprintf(`{"query":"{stopPlace(id:\"NSR:StopPlace:%d\"){id name estimatedCalls(numberOfDepartures:%d){realtime expectedDepartureTime actualDepartureTime destinationDisplay{frontText}serviceJourney{journeyPattern{directionType line{publicCode}}}}}}"}`, stopID, count)
+	query := fmt.Sprintf(`{"query":"{stopPlace(id:\"NSR:StopPlace:%d\"){id name estimatedCalls(numberOfDepartures:%d){realtime expectedDepartureTime actualDepartureTime destinationDisplay{frontText}serviceJourney{operator{id}journeyPattern{directionType line{publicCode}}}}}}"}`, stopID, count)
 	req, err := http.NewRequest("POST", c.URL, strings.NewReader(query))
 	if err != nil {
 		return nil, err
@@ -101,9 +106,13 @@ func parseDepartures(jsonData []byte) ([]Departure, error) {
 	if err := json.Unmarshal(jsonData, &r); err != nil {
 		return nil, err
 	}
+	const operatorPrefix = "ATB:"
 	const timeLayout = "2006-01-02T15:04:05-0700"
 	departures := make([]Departure, 0, len(r.Data.StopPlace.EstimatedCalls))
 	for _, ec := range r.Data.StopPlace.EstimatedCalls {
+		if !strings.HasPrefix(ec.ServiceJourney.Operator.Id, operatorPrefix) {
+			continue // Skip other operators
+		}
 		scheduledDepartureTime, err := time.Parse(timeLayout, ec.ExpectedDepartureTime)
 		if err != nil {
 			return nil, err
diff --git a/entur/testdata/ilsvika.json b/entur/testdata/ilsvika.json
index 2d6c32f..faa8acb 100644
--- a/entur/testdata/ilsvika.json
+++ b/entur/testdata/ilsvika.json
@@ -12,6 +12,9 @@
             "frontText": "Pirbadet via sentrum"
           },
           "serviceJourney": {
+            "operator": {
+              "id": "ATB:Operator:171"
+            },
             "journeyPattern": {
               "directionType": "outbound",
               "line": {
@@ -28,6 +31,9 @@
             "frontText": "Pirbadet via sentrum"
           },
           "serviceJourney": {
+            "operator": {
+              "id": "ATB:Operator:171"
+            },
             "journeyPattern": {
               "directionType": "outbound",
               "line": {
@@ -44,6 +50,9 @@
             "frontText": "Pirbadet via sentrum"
           },
           "serviceJourney": {
+            "operator": {
+              "id": "ATB:Operator:171"
+            },
             "journeyPattern": {
               "directionType": "outbound",
               "line": {
@@ -60,6 +69,9 @@
             "frontText": "Pirbadet via sentrum"
           },
           "serviceJourney": {
+            "operator": {
+              "id": "ATB:Operator:171"
+            },
             "journeyPattern": {
               "directionType": "outbound",
               "line": {
@@ -76,6 +88,9 @@
             "frontText": "Pirbadet via sentrum"
           },
           "serviceJourney": {
+            "operator": {
+              "id": "ATB:Operator:171"
+            },
             "journeyPattern": {
               "directionType": "outbound",
               "line": {
@@ -87,4 +102,4 @@
       ]
     }
   }
-}
\ No newline at end of file
+}
diff --git a/http/http_test.go b/http/http_test.go
index 0304828..dda932e 100644
--- a/http/http_test.go
+++ b/http/http_test.go
@@ -294,6 +294,9 @@ const enturResponse = `{
             "frontText": "Risvollan via sentrum"
           },
           "serviceJourney": {
+            "operator": {
+              "id": "ATB:Operator:171"
+            },
             "journeyPattern": {
               "directionType": "outbound",
               "line": {
@@ -310,6 +313,9 @@ const enturResponse = `{
             "frontText": "Hallset"
           },
           "serviceJourney": {
+            "operator": {
+              "id": "ATB:Operator:171"
+            },
             "journeyPattern": {
               "directionType": "inbound",
               "line": {