Skip to content

Commit

Permalink
completed edge unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@caffix
caffix committed Dec 21, 2024
1 parent cb13bcd commit cb26598
Show file tree
Hide file tree
Showing 2 changed files with 236 additions and 10 deletions.
29 changes: 19 additions & 10 deletions repository/neo4j/edge.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -114,8 +114,11 @@ func (neo *neoRepository) edgeSeen(rel *types.Edge, updated time.Time) error {
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()

query := fmt.Sprintf("MATCH ()-[r]->() WHERE r.elementId = '%s' SET r.updated_at = localDateTime('%s')", rel.ID, timeToNeo4jTime(updated))
_, err := neo4jdb.ExecuteQuery(ctx, neo.db, query, nil,
query := fmt.Sprintf("MATCH ()-[r]->() WHERE elementId(r) = $eid SET r.updated_at = localDateTime('%s')", timeToNeo4jTime(updated))
_, err := neo4jdb.ExecuteQuery(ctx, neo.db, query,
map[string]interface{}{
"eid": rel.ID,
},
neo4jdb.EagerResultTransformer,
neo4jdb.ExecuteQueryWithDatabase(neo.dbname),
)
Expand All @@ -127,7 +130,7 @@ func (neo *neoRepository) FindEdgeById(id string) (*types.Edge, error) {
defer cancel()

result, err := neo4jdb.ExecuteQuery(ctx, neo.db,
"MATCH (from:Entity)-[r]->(to:Entity) WHERE r.elementId = $eid RETURN r, from.entity_id AS fid, to.entity_id AS tid",
"MATCH (from:Entity)-[r]->(to:Entity) WHERE elementId(r) = $eid RETURN r, from.entity_id AS fid, to.entity_id AS tid",
map[string]interface{}{
"eid": id,
},
Expand Down Expand Up @@ -182,12 +185,15 @@ func (neo *neoRepository) IncomingEdges(entity *types.Entity, since time.Time, l
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()

query := fmt.Sprintf("MATCH (:Entity {entity_id: '%s'})<-[r]-(from:Entity) RETURN r, from.entity_id AS fid", entity.ID)
query := "MATCH (:Entity {entity_id: $eid})<-[r]-(from:Entity) RETURN r, from.entity_id AS fid"
if !since.IsZero() {
query = fmt.Sprintf("MATCH (:Entity {entity_id: '%s'})<-[r]-(from:Entity) WHERE r.updated_at >= localDateTime('%s') RETURN r, from.entity_id AS fid", entity.ID, timeToNeo4jTime(since))
query = fmt.Sprintf("MATCH (:Entity {entity_id: $eid})<-[r]-(from:Entity) WHERE r.updated_at >= localDateTime('%s') RETURN r, from.entity_id AS fid", timeToNeo4jTime(since))
}

result, err := neo4jdb.ExecuteQuery(ctx, neo.db, query, nil,
result, err := neo4jdb.ExecuteQuery(ctx, neo.db, query,
map[string]interface{}{
"eid": entity.ID,
},
neo4jdb.EagerResultTransformer,
neo4jdb.ExecuteQueryWithDatabase(neo.dbname),
)
Expand Down Expand Up @@ -250,12 +256,15 @@ func (neo *neoRepository) OutgoingEdges(entity *types.Entity, since time.Time, l
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()

query := fmt.Sprintf("MATCH (:Entity {entity_id: '%s'})-[r]->(to:Entity) RETURN r, to.entity_id AS tid", entity.ID)
query := "MATCH (:Entity {entity_id: $eid})-[r]->(to:Entity) RETURN r, to.entity_id AS tid"
if !since.IsZero() {
query = fmt.Sprintf("MATCH (:Entity {entity_id: '%s'})-[r]->(to:Entity) WHERE r.updated_at >= localDateTime('%s') RETURN r, to.entity_id AS tid", entity.ID, timeToNeo4jTime(since))
query = fmt.Sprintf("MATCH (:Entity {entity_id: $eid})-[r]->(to:Entity) WHERE r.updated_at >= localDateTime('%s') RETURN r, to.entity_id AS tid", timeToNeo4jTime(since))
}

result, err := neo4jdb.ExecuteQuery(ctx, neo.db, query, nil,
result, err := neo4jdb.ExecuteQuery(ctx, neo.db, query,
map[string]interface{}{
"eid": entity.ID,
},
neo4jdb.EagerResultTransformer,
neo4jdb.ExecuteQueryWithDatabase(neo.dbname),
)
Expand Down Expand Up @@ -319,7 +328,7 @@ func (neo *neoRepository) DeleteEdge(id string) error {
defer cancel()

_, err := neo4jdb.ExecuteQuery(ctx, neo.db,
"MATCH ()-[r]->() WHERE r.elementId = $eid DELETE r",
"MATCH ()-[r]->() WHERE elementId(r) = $eid DELETE r",
map[string]interface{}{
"eid": id,
},
Expand Down
217 changes: 217 additions & 0 deletions repository/neo4j/edge_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,3 +3,220 @@
// SPDX-License-Identifier: Apache-2.0

package neo4j

import (
"testing"
"time"

"github.com/owasp-amass/asset-db/types"
"github.com/owasp-amass/open-asset-model/domain"
"github.com/owasp-amass/open-asset-model/relation"
"github.com/stretchr/testify/assert"
)

func TestCreateEdge(t *testing.T) {
from, err := store.CreateEntity(&types.Entity{

Check failure on line 18 in repository/neo4j/edge_test.go

View workflow job for this annotation

GitHub Actions / integration 

undefined: store
Asset: &domain.FQDN{
Name: "create1.edge",
},
})
assert.NoError(t, err)

to, err := store.CreateEntity(&types.Entity{

Check failure on line 25 in repository/neo4j/edge_test.go

View workflow job for this annotation

GitHub Actions / integration 

undefined: store
Asset: &domain.FQDN{
Name: "create2.edge",
},
})
assert.NoError(t, err)

_, err = store.CreateEdge(&types.Edge{

Check failure on line 32 in repository/neo4j/edge_test.go

View workflow job for this annotation

GitHub Actions / integration 

undefined: store
Relation: &relation.SimpleRelation{Name: "invalid_label"},
FromEntity: from,
ToEntity: to,
})
assert.Error(t, err)

first, err := store.CreateEdge(&types.Edge{

Check failure on line 39 in repository/neo4j/edge_test.go

View workflow job for this annotation

GitHub Actions / integration 

undefined: store
Relation: &relation.SimpleRelation{Name: "node"},
FromEntity: from,
ToEntity: to,
})
assert.NoError(t, err)

time.Sleep(250 * time.Millisecond)
second, err := store.CreateEdge(&types.Edge{

Check failure on line 47 in repository/neo4j/edge_test.go

View workflow job for this annotation

GitHub Actions / integration 

undefined: store
Relation: &relation.SimpleRelation{Name: "node"},
FromEntity: from,
ToEntity: to,
})
assert.NoError(t, err)
assert.Equal(t, first.ID, second.ID)
if first.LastSeen == second.LastSeen || second.LastSeen.Before(first.LastSeen) {
t.Errorf("The last seen datetime was not updated")
}
}

func TestFindEdgeById(t *testing.T) {
_, err := store.FindEdgeById("bad_id")

Check failure on line 60 in repository/neo4j/edge_test.go

View workflow job for this annotation

GitHub Actions / integration 

undefined: store
assert.Error(t, err)

from, err := store.CreateEntity(&types.Entity{

Check failure on line 63 in repository/neo4j/edge_test.go

View workflow job for this annotation

GitHub Actions / integration 

undefined: store
Asset: &domain.FQDN{
Name: "find1.edge",
},
})
assert.NoError(t, err)

to, err := store.CreateEntity(&types.Entity{

Check failure on line 70 in repository/neo4j/edge_test.go

View workflow job for this annotation

GitHub Actions / integration 

undefined: store
Asset: &domain.FQDN{
Name: "find2.edge",
},
})
assert.NoError(t, err)

first, err := store.CreateEdge(&types.Edge{

Check failure on line 77 in repository/neo4j/edge_test.go

View workflow job for this annotation

GitHub Actions / integration 

undefined: store
Relation: &relation.SimpleRelation{Name: "node"},
FromEntity: from,
ToEntity: to,
})
assert.NoError(t, err)

second, err := store.FindEdgeById(first.ID)

Check failure on line 84 in repository/neo4j/edge_test.go

View workflow job for this annotation

GitHub Actions / integration 

undefined: store
assert.NoError(t, err)
assert.Equal(t, first.ID, second.ID)
assert.Equal(t, first.FromEntity.ID, second.FromEntity.ID)
assert.Equal(t, first.ToEntity.ID, second.ToEntity.ID)
}

func TestIncomingEdges(t *testing.T) {
from, err := store.CreateEntity(&types.Entity{
Asset: &domain.FQDN{
Name: "incoming1.edge",
},
})
assert.NoError(t, err)

to, err := store.CreateEntity(&types.Entity{
Asset: &domain.FQDN{
Name: "incoming2.edge",
},
})
assert.NoError(t, err)

now := time.Now()
for i := 1; i <= 10; i++ {
created := now.Add(time.Duration(i*-24) * time.Hour)

_, err := store.CreateEdge(&types.Edge{
CreatedAt: created,
LastSeen: created,
Relation: &relation.BasicDNSRelation{
Name: "dns_record",
Header: relation.RRHeader{
RRType: 5,
Class: 0,
TTL: i,
},
},
FromEntity: from,
ToEntity: to,
})
assert.NoError(t, err)
}

_, err = store.IncomingEdges(to, time.Time{}, "invalid_label")
assert.Error(t, err)

edges, err := store.IncomingEdges(to, time.Time{}, "dns_record")
assert.NoError(t, err)
assert.Equal(t, len(edges), 10)

for i := 1; i <= 10; i++ {
since := now.Add(time.Duration(i*-24) * time.Hour)

edges, err := store.IncomingEdges(to, since)
assert.NoError(t, err)
assert.Equal(t, len(edges), i)
}
}

func TestOutgoingEdges(t *testing.T) {
from, err := store.CreateEntity(&types.Entity{
Asset: &domain.FQDN{
Name: "outgoing1.edge",
},
})
assert.NoError(t, err)

to, err := store.CreateEntity(&types.Entity{
Asset: &domain.FQDN{
Name: "outgoing2.edge",
},
})
assert.NoError(t, err)

now := time.Now()
for i := 1; i <= 10; i++ {
created := now.Add(time.Duration(i*-24) * time.Hour)

_, err := store.CreateEdge(&types.Edge{
CreatedAt: created,
LastSeen: created,
Relation: &relation.BasicDNSRelation{
Name: "dns_record",
Header: relation.RRHeader{
RRType: 5,
Class: 0,
TTL: i,
},
},
FromEntity: from,
ToEntity: to,
})
assert.NoError(t, err)
}

_, err = store.OutgoingEdges(from, time.Time{}, "invalid_label")
assert.Error(t, err)

edges, err := store.OutgoingEdges(from, time.Time{}, "dns_record")
assert.NoError(t, err)
assert.Equal(t, len(edges), 10)

for i := 1; i <= 10; i++ {
since := now.Add(time.Duration(i*-24) * time.Hour)

edges, err := store.OutgoingEdges(from, since)
assert.NoError(t, err)
assert.Equal(t, len(edges), i)
}
}

func TestDeleteEdge(t *testing.T) {
from, err := store.CreateEntity(&types.Entity{
Asset: &domain.FQDN{
Name: "delete1.edge",
},
})
assert.NoError(t, err)

to, err := store.CreateEntity(&types.Entity{
Asset: &domain.FQDN{
Name: "delete2.edge",
},
})
assert.NoError(t, err)

edge, err := store.CreateEdge(&types.Edge{
Relation: &relation.SimpleRelation{Name: "node"},
FromEntity: from,
ToEntity: to,
})
assert.NoError(t, err)

err = store.DeleteEdge(edge.ID)
assert.NoError(t, err)

_, err = store.FindEdgeById(edge.ID)
assert.Error(t, err)
}

0 comments on commit cb26598

Please sign in to comment.